# Training an LLM Email Triage Agent with GRPO Reinforcement Learning

> **TL;DR**: We built an OpenEnv-based RL environment that teaches a 3B LLM to 
> automatically triage a corporate inbox — classifying urgency, routing to the 
> right team, and drafting replies. Using GRPO, the agent improved from an avg 
> reward of 3.47 → 4.10+ over 200 training steps, with correct routing jumping 
> from ~40% to ~85%.

---

## The Problem

Corporate inboxes are chaos. Every day, support teams manually read hundreds 
of emails — figuring out if something is urgent, who should handle it, and 
what to reply. A hacked account alert needs immediate escalation. A lottery 
spam needs to go straight to trash. An overdue invoice needs the billing team. 
A sales lead needs the sales team.

Getting this wrong is expensive: SLA breaches, missed leads, frustrated 
customers. Getting it right manually is time-consuming and error-prone.

We built an RL environment to teach an LLM to do this automatically — and 
then actually trained it.

---

## What We Built

An **OpenEnv-based email triage environment** where an LLM agent receives a 
corporate inbox and must triage each email by:

- **Classifying** it — spam? billing inquiry? security incident? sales lead?
- **Prioritizing** it — urgent / high / medium / low / spam
- **Routing** it — to billing, sales, support tier 1/2, HR, trash, archive...
- **Drafting** a professional reply to the sender

### What Makes It Interesting

The environment has real constraints that force the agent to reason carefully:

**SLA Warnings** — urgent emails must be handled within 2 steps or you get 
penalized. The agent sees active SLA warnings in its observation and must 
prioritize accordingly.

**Escalation Budget** — only 3 emails per session can be flagged for human 
review. Flag spam or low-priority emails and you waste your budget, getting 
penalized on legitimate escalations later. The agent must learn to be 
selective.

**Sequential Decisions** — 5 emails per episode, processed one at a time. 
Each decision affects the remaining state (queue capacity, escalation budget).

**Multi-dimensional Reward** — the environment scores each action across 
priority accuracy, category correctness, routing quality, reply draft quality 
(20% of score), and SLA compliance. There's no single shortcut to game.

---

## Environment Design

```
Observation:
  current_email     → the email to triage now
  inbox             → all emails in the episode
  escalation_budget → how many flag_review actions remain
  sla_warnings      → emails approaching deadline
  team_queues       → routing capacity per destination

Action:
  email_id          → which email (must match exactly)
  priority          → urgent | high | medium | low | spam
  category          → customer_complaint | billing_inquiry | 
                       technical_support | sales_lead | internal_hr |
                       legal_compliance | spam_phishing | general_inquiry
  route_to          → support_tier1 | support_tier2 | billing | sales |
                       legal | hr | management | trash | archive
  summary           → ≤280 char description
  flag_review       → true/false (uses escalation budget)
  reply_draft       → professional reply to sender

Reward:
  Per email: priority (~0.2) + category (~0.2) + routing (~0.3) 
           + reply quality (~0.2) + escalation appropriateness (~0.1)
  Penalties: SLA breach, invalid field values, budget overuse
  Max score: ~4.5 per episode (5 emails × ~0.9 each)
```

---

## Training with GRPO

We used **GRPO (Group Relative Policy Optimization)** via HuggingFace TRL 
to train a `Qwen/Qwen2.5-3B-Instruct` model with LoRA adapters.

### Why GRPO?

GRPO is well-suited for this task because:
- It doesn't need a learned reward model — the environment IS the verifier
- It compares episode pairs and reinforces better decisions relatively
- It's more compute-efficient than PPO for small models

### Why LoRA?

We only update ~2M out of 3B parameters using LoRA (r=4, targeting q_proj 
and v_proj). The base model already knows language, reasoning, and JSON 
formatting from pretraining on 18T tokens. LoRA just steers that existing 
knowledge toward email triage specifically.

### Training Configuration

```python
MODEL_NAME      = "Qwen/Qwen2.5-3B-Instruct"
TRAINING_STEPS  = 200
BATCH_SIZE      = 1
NUM_GENERATIONS = 4        # episodes compared per GRPO step
learning_rate   = 2e-6
lora_r          = 4
lora_alpha      = 8
target_modules  = ["q_proj", "v_proj"]
lr_scheduler    = "cosine"
warmup_steps    = 15
```

### The Training Loop

```
For each step:
  1. Run 4 episodes through the email environment
  2. Environment scores each decision (priority + category + routing + reply)
  3. GRPO compares: which episodes scored higher?
  4. Advantage = (reward - mean) / std per group
  5. Update LoRA weights: reinforce high-reward decisions,
     discourage low-reward ones
  6. Repeat 200 times
```

---

## Key Engineering Decisions

### 1. Chat Template is Non-Negotiable

Using `apply_chat_template()` was the single biggest fix in the entire 
project. Without it, even a 7B instruct model produces garbage — the model 
expects a specific `<|im_start|>system\n...<|im_end|>` format that was baked 
in during instruction fine-tuning. Sending raw text completely bypasses this.

### 2. Validation as a Safety Net

We built a `validate_action()` function that catches and corrects invalid 
outputs before they reach the environment:

```python
# Catches: invented categories like "sales_leads", "spammer_alert"
# Catches: array-wrapped JSON [{...}] instead of objects {...}
# Catches: null reply_drafts
# Catches: route/category mismatches (e.g. spam → sales)
# Catches: escalation budget abuse (flagging spam for human review)
# Infers:  correct values from email body keywords when model fails
```

This lifted the reward floor from 2.78 to 3.60 — the most impactful single 
change after the chat template fix.

### 3. Escalation Budget Management

The model initially flagged spam and low-priority emails for human review, 
exhausting the escalation budget early and crashing episode rewards to ~0.45. 
We added explicit budget tracking and a rule: only flag `urgent` or `high` 
priority non-spam emails. This eliminated the reward crashes entirely.

### 4. model.eval() During Rollout

GRPO trains the model while simultaneously using it for inference. Without 
explicitly calling `model.eval()` during rollouts and `model.train()` after, 
batch normalization and dropout behave incorrectly, causing unstable 
generation during training steps.

---

## Results

### Reward Curve

![Reward Curve](training/plots/reward_curve.png)

The agent improved from **3.47 → 4.10+** average reward over 200 steps — 
an 18% improvement. More importantly, the floor lifted from 2.78 to 3.60, 
meaning the agent became consistently reliable, not just occasionally good.

### Key Metrics

| Metric | Baseline | After Training |
|--------|----------|----------------|
| Avg episode reward | 3.47 | 4.10+ |
| Best episode reward | 3.65 | 4.41 |
| Worst episode reward | 2.78 | 3.60 |
| Final evaluation score | — | **4.189 / 5.0** |
| Correct routing | ~40% | **100% (5/5)** |
| Null reply drafts | frequent | eliminated |
| Invalid categories | occasional | eliminated |
| Escalation budget abuse | frequent | eliminated |

### Before vs After — Routing Comparison (Final Evaluation)

| Email | Untrained Agent | Trained Agent | Reward |
|-------|-----------------|---------------|--------|
| 🔴 Account hacked | general_inquiry → support_tier1 | customer_complaint → support_tier2 ✅ | 0.705 |
| 🗑️ Lottery spam | urgent → support_tier1 | spam_phishing → trash ✅ | 0.840 |
| 👥 Team lunch | low → trash | internal_hr → hr ✅ | 0.988 |
| 💰 Overdue invoice | medium → support_tier1 | billing_inquiry → billing ✅ | 0.798 |
| 📈 Enterprise lead | general_inquiry → archive | sales_lead → sales ✅ | 0.858 |
| **Total** | | | **4.189 / ~5.0** |

---

## Key Lessons Learned

**1. Model size matters more than you think for structured output.**
The 0.5B model produced pure garbage (LaTeX, Chinese exam questions, C++ 
code). The 1.5B was marginal. The 3B consistently produced valid JSON from 
step 1. For tasks requiring strict output format, don't go below 3B.

**2. The reward floor matters as much as the average.**
A model that scores 4.0 on average but occasionally crashes to 0.4 is less 
useful than one that scores 3.8 consistently. Most of our engineering effort 
went into lifting the floor, not the ceiling.

**3. RL on small fixed datasets still works.**
With only 5 fixed emails, we worried the model would memorize rather than 
generalize. But GRPO's relative reward signal still found meaningful 
gradients. The model learned general routing principles — spam goes to trash, 
billing goes to billing — not just specific email answers.

**4. Validation beats prompting for reliability.**
No matter how carefully you word the system prompt, a 3B model will 
occasionally output `sales_leads` instead of `sales_lead`, or wrap JSON in 
an array. A deterministic validation function that catches and corrects these 
errors is far more reliable than trying to prompt your way out of them.

**5. Environment constraints create richer learning signals.**
The escalation budget and SLA warnings force the model to reason about 
trade-offs, not just classify emails in isolation. This is what makes the 
environment genuinely interesting from an RL perspective — the agent has to 
learn a policy, not just a lookup table.

---

## Try It Yourself

**Environment**: [HF Space](https://huggingface.co/spaces/Vansh04092003-multi-agent-email-env2)

**Training notebook**: [Open in Colab](#) ← add your link

**wandb run**: [Live metrics](#) ← add your link

```python
import requests

BASE_URL = "https://Vansh04092003-multi-agent-email-env2.hf.space"

# Start a fresh episode
obs = requests.post(f"{BASE_URL}/reset").json()["observation"]
print(f"First email: {obs['current_email']['header']['subject']}")

# Submit a triage decision
action = {
    "email_id":    "e001",
    "priority":    "urgent",
    "category":    "customer_complaint",
    "route_to":    "support_tier2",
    "summary":     "Account hacked, needs immediate lock.",
    "flag_review": True,
    "reply_draft": "We are securing your account immediately."
}
result = requests.post(f"{BASE_URL}/step", json=action).json()
print(f"Reward: {result['reward']:.3f}")
```

---

## What's Next

- **More email diversity**: The current environment has 5 fixed emails. 
  Adding procedurally generated emails would make the learned policy more 
  robust and generalizable.
- **Larger model**: Upgrading to Qwen2.5-7B or Qwen3-8B would likely push 
  the reward ceiling from 4.4 toward 4.8+.
- **Process-level rewards**: Currently rewards are per-email. Adding 
  step-level rewards for good reasoning chains (thinking tokens) could 
  improve sample efficiency.
- **Multi-agent extension**: Multiple agents collaborating on inbox triage, 
  with handoffs between support tiers — a natural extension of the current 
  environment toward Theme #1.

---

*Built for the OpenEnv Hackathon India 2026 — Theme #3.2 Personalized Tasks.*

*Environment: [Vansh04092003/multi-agent-email-env2](https://huggingface.co/spaces/Vansh04092003-multi-agent-email-env2)*