GRPO-Trained Cost-Aware Multi-Tool Routing Agent

Research question: can GRPO (Group Relative Policy Optimization) train a small language model to route between tools of different cost — SQL queries, document/web retrieval, a calculator — choosing the cheapest sufficient tool combination while still answering correctly? And separately: does whatever it learns generalize to a tool combination it never saw during training, or does it just memorize which templates map to which tools?

This is a portfolio artifact built to be honest and reproducible first, impressive second. Where the real result is modest or partially negative, that's reported plainly — see docs/findings.md for the full, undecorated account, including a headline finding that is a genuine generalization failure, not a success story.

What's actually in here

• A synthetic company world (SQLite DB + internal doc corpus + external web corpus, deterministically generated, structurally disjoint — no fact is answerable from more than one source) and ~4,230 templated tasks over it. Templated, not natural language — see "Honest limitations" below for why that tradeoff was deliberate.
• 4 tools with distinct costs (sql_query=1, doc_search=2, web_search=5, calculator=0.5), a strict fenced-JSON action format, and a programmatic (never LLM-judged) reward combining correctness, cost-efficiency, and a format-validity penalty.
• 6 baselines (always-cheapest, always-all-tools, a static-confidence-threshold heuristic mirroring a hardcoded production system, an oracle upper bound, and the zero-shot base model) plus the actual GRPO-trained policy, all scored through the same evaluation harness.
• A real GRPO training run — Qwen2.5-1.5B-Instruct, 4-bit NF4 + LoRA, on an actual GPU (Colab T4), with full diagnostics (reward, entropy, KL, gradient norm, distinct-trajectory diversity, per-tool latency) logged every step.

Quickstart

uv venv --python 3.11 .venv && uv pip install --python .venv -r requirements.txt

# Part 1: build the world + task splits (deterministic, seeded)
python scripts/build_dataset.py

# Part 3/9: unit tests (129, no GPU needed)
pytest tests/ -q

# Part 4: reward profiling (the untrained model, mandatory gate before training)
python scripts/profile_rewards.py --smoke-test              # plumbing check, CPU, minutes
python scripts/profile_rewards.py --model-id Qwen/Qwen2.5-1.5B-Instruct --n-tasks 150 --k 8  # real, needs GPU

# Part 5: baselines (fast, no model needed except the optional zero-shot one)
python scripts/run_baselines.py
python scripts/run_baselines.py --budget-mode hard

# Part 6: GRPO training (smoke test is CPU-feasible; the real run needs a GPU)
python scripts/train_grpo.py --smoke-test
python scripts/train_grpo.py --max-steps 100                # real run, see docs/findings.md for why capped

# Part 7: evaluation + the final results table/plots
python scripts/evaluate.py --methods zero_shot grpo_trained --pass-k 1 4 --n-tasks 60
python scripts/build_results_table.py
python scripts/make_training_plots.py
python scripts/dump_trajectories.py

Everything through run_baselines.py runs comfortably on a CPU-only machine. The real profiling/training/evaluation runs need a real GPU (16GB, e.g. a free-tier Colab T4) — see docs/colab_guide.md for exact setup steps, including the resume/backup patterns that turned out to matter a lot in practice (see findings).

Results

Full table: results/results_table.md / .csv. Headline rows (60-task eval subset, pass@4 sampling, soft budget mode unless noted):

method	split	success_rate	mean_cost_per_solved	success_per_unit_cost
zero_shot	test	0.708	1.012	0.334
grpo_trained	test	0.762	1.000	0.421
zero_shot	test_heldout_combo	0.0	—	—
grpo_trained	test_heldout_combo	0.0	—	—
oracle (upper bound)	test_heldout_combo	0.5	7.500	0.069

GRPO training improved the in-distribution policy — higher success rate and lower cost per solved task than the untrained model, the intended cost-aware-routing behavior. It did not generalize to the held-out tool combination at all — 0% for both models, 240/240 attempts each. Only oracle (told exactly which tools to use) gets any traction there. This gap between test and test_heldout_combo is the single most important number in the project — see docs/findings.md for the full analysis, including why (a 100-step, compute-capped run — see below) and what that does and doesn't tell us.

Also notable: the static_threshold baseline — a hardcoded retrieval-confidence-escalation heuristic modeling a real production pattern — scores 0% everywhere, because it structurally never calls sql_query or calculator. This is the specific kind of brittleness a learned router is meant to fix; that our own learned router also failed to generalize (above) is a pointed, honest parallel rather than a clean redemption story.

Plots: training curves and the mode-collapse/optimization-health diagnostics (results/training_*.png), the real reward-profile distribution (results/reward_profile.png), cost-vs-success and latency-vs-success scatter plots with a Pareto frontier (results/cost_vs_success.png, results/latency_vs_success.png), and the test-vs-heldout-combo comparison (results/test_vs_heldout_bars.png). Ten qualitative trajectories (5 success, 5 failure, each with a diagnosis) are in results/trajectories/.

Honest limitations

• Tasks are templated, not natural language. ~18 conceptual templates with a few phrasing variants each, parametrized over the generated world's actual rows. This was a deliberate tradeoff, not an oversight: the entire point of this project is a reward computed by exact programmatic checks, with no LLM-as-judge anywhere — that requires knowing precisely what was asked and what the one correct answer is, which templated generation guarantees and free-form natural language does not.
• The real training run was capped at 100 steps, not a full epoch. Episode generation is sequential (not batched across the GPU), making a full epoch ~8 days of continuous compute — infeasible for this project's compute budget. See docs/findings.md for the full accounting and what batching the generation loop would change.
• The 60-task evaluation subset for zero-shot/GRPO-trained was not stratified, and happened to land entirely on difficulty=1 tasks for test and difficulty=2 for test_heldout_combo — the difficulty/multihop breakdown (results/results_by_difficulty.md) is only informative for the baselines (evaluated at full scale: 520/200 tasks), not the model-based methods.
• KL regularization was disabled (beta=0.0) for the real training run specifically to fit in a 15GB GPU's memory — a real, disclosed tradeoff, not a silent one.
• The SFT baseline was not run (explicitly optional in the project spec) — without it, the in-distribution gain can't be cleanly attributed to RL versus what plain imitation of oracle trajectories would have achieved for free.
• zero_shot/grpo_trained were only evaluated in soft budget mode — the hard-vs-soft budget comparison (a real, interesting finding — see docs/findings.md) was only run for the cheap scripted baselines due to compute constraints.

None of these are hidden — each is discussed with its actual consequence in docs/findings.md, which is written as an honest lab notebook, not as marketing for the project.

Repository layout

configs/     env.yaml, model.yaml, train_grpo.yaml — every numeric knob, one source of truth
src/
  env/       world generation, tools, tasks, RoutingEnv (Gym-style, model-free)
  agent/     action format, parsing, rollout loop, HF/LoRA policy wrappers
  rewards/   reward function + profiling logic (both pure, unit-tested)
  baselines/ the 4 scripted baselines + rule-based query synthesis
  train/     GRPO training orchestration, diagnostics, tokenization/masking
  eval/      evaluation harness, pass@k, plots, results-table building
scripts/     one CLI entry point per pipeline stage (see Quickstart)
tests/       129 tests, deterministic, no GPU needed
results/     every real number and plot referenced above
docs/        reward_design.md, findings.md, colab_guide.md