Build Hour: Agent RFT

Optimizing Agent Performance

Improving the performance of an AI agent, a model that interacts with the outside world via tools, involves a hierarchy of optimization techniques. While prompt engineering is a great starting point, further gains can be achieved by refining the task itself—simplifying it, adding guardrails, or improving the tools’ descriptions and functionality.

However, when these methods reach their limits, fine-tuning becomes the next step to achieve peak performance. Agent Reinforcement Fine-Tuning (Agent RFT) is a method that trains the agent end-to-end on a specific task, fundamentally changing the model’s weights to teach it what constitutes good and bad behavior.

Understanding Agent RFT

Agent RFT enhances the existing Reinforcement Fine-Tuning (RFT) framework by specifically enabling agents to use tools during the training process. This is a crucial distinction from Base RFT. During training, the agent explores various ways of calling your unique tools, learning from the outcomes to improve its strategy over time.

Two major additions make this possible:

1. Tool Calls via Endpoints: The model can call your custom tools hosted on external endpoints during the training process.
2. Custom Grader via Endpoint: You can provide a custom reward signal through an endpoint, allowing you to define what success looks like in a way that is specific to your business logic and task.

This marks the first time OpenAI models can interact with the outside world during training, both through tool use and by receiving feedback from your custom reward signal. This process helps close the “domain gap” between the general knowledge of the base model and the specific context of your business and tools, leading to several key benefits:

• Improved Performance: The model learns to reason more effectively across tool outputs and use the tools themselves more efficiently.
• Lower Latency: By learning the most efficient path to a solution, the model often uses fewer tool calls and reasoning tokens, resulting in faster response times for the end-user. This happens naturally due to a light penalty on token usage, but can also be enforced by setting a tool-call budget.
• Sample Efficiency: Agent RFT can be highly effective even with relatively small, high-quality datasets, which is critical in domains where training data is scarce.

Practical Demonstration: FinQA Benchmark

To illustrate Agent RFT, the team demonstrated fine-tuning an agent on a modified version of the FinQA (Financial Question Answering) benchmark.

• The Task: The agent was given a financial question without the relevant report and was required to find the correct information within a corpus of 2,800 financial reports to answer it. The task was further constrained to a maximum of 10 tool calls.
• Tools Provided:
  • search: A semantic search tool.
  • list: A tool to list directories and file paths.
  • cat: A tool to return the content of a document given a path.
• The Grader: A model-based grader (GPT-4o) was used instead of a simple string match. This provided more robust evaluation, allowing for partial credit (e.g., for rounding errors) and flexibility in answer formatting, ensuring the agent received a meaningful learning signal.

The results were significant. After just 10 training steps, the fine-tuned model showed:

• An 11 percentage point increase in reward (task success).
• A 10% reduction in average latency.
• A 40% reduction in the average number of tool calls (from 6.9 to 4.2), indicating the model learned a much more efficient policy for finding and reasoning with information.

Customer Spotlight: Cognition on Improving Devin

Sampriti Panda, a Research Engineer at Cognition, shared how they used Agent RFT to improve Devin, their autonomous AI engineer.

• The Goal: To make Devin’s initial “planning mode” faster and more accurate. The sub-task was for the agent to identify the correct files in a codebase that need to be edited to solve a user’s request.
• The Setup: The agent was given read file and shell tools (for grep, find, etc.). The reward signal was the F1 score, balancing precision and recall of the identified files.
• The Results: Agent RFT led to a massive performance boost. The fine-tuned model learned to use parallel tool calls effectively, reducing the number of back-and-forth interactions required to complete the planning phase from 8-10 down to just 4.
• Infrastructure Insights: Cognition ran each training rollout in an isolated VM to prevent destructive actions (e.g., rm -rf). They noted the training platform’s traffic is “bursty” and emphasized the importance of robust monitoring, as infrastructure failures can result in zero-reward signals that negatively impact the training process.

Additional Success Stories

• Ambience: Used Agent RFT for an ICD-10 medical coding agent, achieving a significant F1 score improvement and an 18% reduction in latency.
• Rogo: Built a financial reasoning agent. Using a custom LLM-as-a-judge grader, they saw a 21% increase in core ML performance and learned the importance of creating a “watertight” grader to prevent the model from “reward hacking” or finding exploits.
• Mako: Developed a GPU kernel writing agent for new hardware, a domain with very little training data. Agent RFT beat the state-of-the-art by 72% with as few as 100 prompts.

Best Practices and Recommended Workflow

To succeed with Agent RFT, focus on:

1. Task & Data:
   • Ensure your task is well-specified and constrained.
   • Start with a model that has non-zero baseline performance. The model needs some initial success to learn from.
   • Prioritize a high-quality dataset over a large one; 150-1000 samples can be very effective.
2. Infrastructure:
   • Mirror your production environment in your tool endpoints.
   • Invest heavily in designing a robust grader that is hard to game and provides a gradient of scores (not just binary 0 or 1).
   • Limit the output length of tool calls to keep training efficient.

The recommended workflow is:

1. Build a high-quality dataset that matches production traffic.
2. Establish a baseline performance with a base model like GPT-4o.
3. Attempt to optimize through other means first (prompting, tool improvements).
4. Once those are exhausted, turn to Agent RFT to fundamentally alter the model’s weights and achieve the next level of performance.

• Tags:
• Fine-Tuning
• Reinforcement Learning
• Devin
• Tool Use
• LLM Agents
• Agent RFT
• Cognition