Reinforcement Learning for MCP Tool Calling

SFT Approach:
Prompt -> Golden Action -> Golden Observation -> Golden Answer
(Zero self-correction capability on edge cases)

GRPO Approach:
Prompt -> Generate N Parallel Trajectories {T_1, T_2, ... T_N}
       -> Execute tools against live MCP sandbox
       -> Calculate Verifiable Reward (Tests pass? SQL valid?)
       -> Optimize policy weights relative to group baseline

{
  "system_prompt": "You are a Model Context Protocol reasoning agent. You have access to the following MCP tools:\n- name: query_database\n  inputSchema: {\"sql\": {\"type\": \"string\"}}\n- name: run_linter\n  inputSchema: {\"path\": {\"type\": \"string\"}}",
  "task": "Find all inactive customer records in postgres and verify schema integrity.",
  "ground_truth_assertions": [
    "SELECT COUNT(*) FROM customers WHERE status = 'inactive'",
    "exit_code == 0"
  ]
}

# mcp_rl_reward.py
import json
import re

def compute_mcp_reward(completions, ground_truth, **kwargs) -> list[float]:
    """
    Reward function scoring MCP tool calling trajectories:
    1. Schema Syntax: Valid JSON-RPC structure (+0.3)
    2. Correct Tool Selection (+0.3)
    3. Correct Execution Outcome (+0.4)
    4. Penalty for Hallucinated Tools (-0.5)
    """
    rewards = []

    for text in completions:
        score = 0.0

        # Check for tool call tags <tool_call>{...}</tool_call>
        match = re.search(r"<tool_call>(.*?)</tool_call>", text, re.DOTALL)
        if not match:
            rewards.append(0.0)
            continue

        raw_json = match.group(1).strip()
        try:
            call_obj = json.loads(raw_json)
        except json.JSONDecodeError:
            rewards.append(0.1)  # Partial credit for attempting format
            continue

        score += 0.3  # Valid JSON schema syntax

        tool_name = call_obj.get("name")
        args = call_obj.get("arguments", {})

        # Verify against allowed tools catalog
        if tool_name not in ["query_database", "run_linter"]:
            rewards.append(-0.5)  # Penalty for hallucination
            continue

        score += 0.3  # Correct known tool

        # Verify task-specific arguments
        if tool_name == "query_database" and "inactive" in args.get("sql", "").lower():
            score += 0.4  # Successfully formulated the expected query

        rewards.append(score)

    return rewards

# train_mcp_grpo.py
from trl import GRPOTrainer, GRPOConfig
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen2.5-Coder-7B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")

training_args = GRPOConfig(
    output_dir="./qwen-mcp-grpo-checkpoints",
    learning_rate=1e-5,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    num_generations=8,  # Sample 8 trajectories per prompt
    max_prompt_length=1024,
    max_completion_length=1024,
    logging_steps=10
)

# trainer = GRPOTrainer(
#     model=model,
#     reward_funcs=compute_mcp_reward,
#     args=training_args,
#     train_dataset=dataset
# )
# trainer.train()