A Practical Guide to Agent Reinforcement Fine-Tuning

# step1_prior_check.py
import json, requests, random
from collections import Counter

BASE_MODEL_URL = "http://localhost:8001/v1/chat/completions"   # vLLM endpoint
N_ROLLOUTS = 5

def run_agent(prompt, tools_spec):
    """Send the prompt to the base model with tools; return (success: bool, n_tool_calls: int)."""
    # ... your model + tool-calling loop here ...
    # Returns whether the final answer matched ground truth.
    pass

samples = [json.loads(line) for line in open("data/research_val.jsonl")]
results = []
for s in samples[:50]:
    successes = sum(run_agent(s["prompt"], s["tools"])[0] for _ in range(N_ROLLOUTS))
    rate = successes / N_ROLLOUTS
    results.append((s["id"], rate))

# Bucket each sample as never / sometimes / always
buckets = Counter(
    "always"    if r == 1.0 else
    "sometimes" if 0 < r < 1 else
    "never"
    for _, r in results
)
print(buckets)

{
  "id": "row-2842",
  "prompt": "What was Henderson Industries' Q3 2025 operating margin?",
  "reference_answer": "13.4%",
  "expected": {
    "tolerance_rel": 0.02,
    "evidence_path": "filings/henderson-10q-2025-q3.html",
    "must_use_tools": ["search", "read"]
  },
  "tools": [
    {"type": "function", "function": {"name": "search", "parameters": {...}}},
    {"type": "function", "function": {"name": "list",   "parameters": {...}}},
    {"type": "function", "function": {"name": "read",   "parameters": {...}}}
  ],
  "metadata": {"task_type": "numerical_qa", "difficulty": "medium"}
}

# step2_validate_dataset.py
import json
from pathlib import Path

REQUIRED = {"id", "prompt", "reference_answer"}
ROWS = []
for path in ["data/research_train.jsonl", "data/research_val.jsonl"]:
    with open(path) as f:
        for ln, line in enumerate(f, 1):
            row = json.loads(line)
            missing = REQUIRED - row.keys()
            assert not missing, f"{path}:{ln} missing {missing}"
            assert row["id"] not in {r["id"] for r in ROWS}, f"dup id {row['id']}"
            ROWS.append(row)

print(f"Total rows: {len(ROWS)}")
print(f"Train/val: {sum(1 for r in ROWS if r['id'].startswith('train'))} / "
      f"{sum(1 for r in ROWS if r['id'].startswith('val'))}")
print(f"Difficulty distribution: "
      f"{dict(Counter(r['metadata']['difficulty'] for r in ROWS))}")

# step3_tool_server.py
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.security import HTTPBearer

app = FastAPI()
security = HTTPBearer()

def auth(creds=Depends(security)):
    if creds.credentials != os.environ["TOOL_BEARER"]:
        raise HTTPException(401, "bad token")

@app.post("/search")
async def search(req: Request, _=Depends(auth)):
    body = await req.json()
    args = body.get("arguments", {})
    call_id = body.get("call_id")
    query = args.get("query", "").strip()
    if not query:
        # Structured error, not 500 — see Watch For below
        return {"output": "Error: query is required", "call_id": call_id}
    hits = embed_and_search(query, top_k=3)         # your retrieval logic
    return {"output": format_hits(hits), "call_id": call_id}

@app.post("/list")
async def list_paths(req: Request, _=Depends(auth)):
    body = await req.json()
    prefix = body.get("arguments", {}).get("prefix", "")
    paths = corpus_index.list(prefix)
    return {"output": "\n".join(paths), "call_id": body.get("call_id")}

@app.post("/read")
async def read(req: Request, _=Depends(auth)):
    body = await req.json()
    path = body.get("arguments", {}).get("path")
    return {"output": corpus_index.read(path)[:8000],   # token budget!
            "call_id": body.get("call_id")}

$ curl -s -X POST https://tools.example.com/search \
    -H "Authorization: Bearer $TOOL_BEARER" \
    -H "Content-Type: application/json" \
    -d '{"arguments": {"query": "Henderson Q3 operating margin"},
          "call_id": "test-001"}' | jq

# step4_grader.py
GRADER = {
    "type": "score_model",
    "name": "numerical_grader",
    "model": "gpt-4.1",
    "range": [0, 1],
    "pass_threshold": 0.75,
    "sampling_params": {"temperature": 0},
    "input": [
        {"role": "system", "content": """\
You will be given a Reference Answer and a Model Answer.
Score the Model Answer on a 0-1 scale.

Return 1 if BOTH:
- the Model Answer contains only the final numeric answer
- the value matches the Reference Answer up to format

Format variations that still count as correct:
- currency symbols ($, USD, €)
- magnitude suffixes (M, million, K, B)
- percent forms (7% vs 0.07)
- commas and whitespace

Return 0.5 if off by a tenth of a percent or less, or a true rounding error.
Return 0 otherwise. Reply with ONLY the numeric score."""},
        {"role": "user", "content":
            "- Reference Answer: {{item.reference_answer}}\n"
            "- Model Answer: {{sample.output_text}}"}
    ],
}

# step4_grader_test.py
cases = [
    ({"reference_answer": "13.4%"}, {"output_text": "13.4%"}),       # exact
    ({"reference_answer": "13.4%"}, {"output_text": "0.134"}),       # format variant
    ({"reference_answer": "13.4%"}, {"output_text": "13.41%"}),      # rounding
    ({"reference_answer": "13.4%"}, {"output_text": "12.0%"}),       # wrong
]
for item, sample in cases:
    print(f"ref={item['reference_answer']:6}  out={sample['output_text']:8}  "
          f"reward={run_grader(item, sample):.2f}")

# step5_config.yaml
algorithm: grpo            # group-relative policy optimization

# The four numbers that matter most
group_size: 8              # rollouts per prompt; below 4 the baseline is too noisy
kl_coef: 0.01              # β: KL penalty against the frozen reference. Tune carefully.
learning_rate: 1.0e-5      # 10-100× lower than SFT
epochs: 2                  # agent RFT overfits fast on small datasets

# Mostly-fine-as-is
batch_size: 16             # max within memory after rollout buffer
max_trajectory_tokens: 4096
max_tool_calls: 12         # rollout tool-call budget
warmup_steps: 5
seed: 42

# Inputs
train_path: data/research_train.jsonl
val_path:   data/research_val.jsonl
base_model: Qwen/Qwen2.5-7B-Instruct
lora:
  rank: 32
  alpha: 64
  target_modules: [q_proj, k_proj, v_proj, o_proj]

# step6_train.py
from trl import GRPOTrainer, GRPOConfig
from peft import LoraConfig
from datasets import load_dataset
from utils import build_grader, build_tools

config = GRPOConfig(
    learning_rate=1e-5,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    num_train_epochs=2,
    num_generations=8,                  # group size
    beta=0.01,                           # KL coefficient
    max_completion_length=4096,
    bf16=True,
    logging_steps=1,
    output_dir="./runs/research-rft-v1",
)

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-7B-Instruct",
    reward_funcs=[build_grader(GRADER)],   # from Step 4
    train_dataset=load_dataset("json", data_files="data/research_train.jsonl")["train"],
    args=config,
    peft_config=LoraConfig(r=32, lora_alpha=64,
                           target_modules=["q_proj","k_proj","v_proj","o_proj"]),
)
trainer.train()

# step7_evaluate.py
import json, statistics

def evaluate(model_id, val_path):
    out = []
    for row in (json.loads(l) for l in open(val_path)):
        sample = run_agent(model_id, row["prompt"], row["tools"])
        reward = run_grader(row, sample)
        out.append({"id": row["id"],
                    "reward": reward,
                    "tool_calls": sample["n_tool_calls"]})
    return out

base = evaluate("Qwen/Qwen2.5-7B-Instruct", "data/research_val.jsonl")
tuned = evaluate("./runs/research-rft-v1/checkpoint-150", "data/research_val.jsonl")

# Per-sample diff
diffs = []
for b_, t_ in zip(base, tuned):
    diffs.append({
        "id": b_["id"],
        "d_reward": t_["reward"] - b_["reward"],
        "d_tools":  t_["tool_calls"] - b_["tool_calls"],
    })

print(f"mean Δreward: {statistics.mean(d['d_reward'] for d in diffs):+.3f}")
print(f"mean Δtools:  {statistics.mean(d['d_tools']  for d in diffs):+.2f}")
print(f"improved:    {sum(1 for d in diffs if d['d_reward'] > 0)}/{len(diffs)}")
print(f"regressed:   {sum(1 for d in diffs if d['d_reward'] < 0)}/{len(diffs)}")
print(f"unchanged:   {sum(1 for d in diffs if d['d_reward'] == 0)}/{len(diffs)}")

# step8_inspect.py
import json
from pathlib import Path

# Sort regressions worst-first
regressions = sorted([d for d in diffs if d["d_reward"] < 0],
                     key=lambda d: d["d_reward"])

for d in regressions[:5]:
    print(f"\n=== {d['id']}  Δreward={d['d_reward']:+.2f}  Δtools={d['d_tools']:+}")
    row = next(r for r in val_rows if r["id"] == d["id"])
    print(f"PROMPT:    {row['prompt']}")
    print(f"REFERENCE: {row['reference_answer']}")

    trace = load_trace(model="tuned", row_id=d["id"])
    for turn in trace["turns"]:
        if turn["type"] == "tool_call":
            print(f"  → {turn['name']}({turn['args']})")
        elif turn["type"] == "tool_result":
            print(f"  ← {turn['output'][:80]}")
        elif turn["type"] == "message":
            print(f"  [model] {turn['content'][:120]}")