ProtAgent: Post-Training Compact Protein Tool Agents with Interaction-Level Credit Assignment

Abstract

Protein annotation agents must choose suitable bioinformatics tools, interpret their outputs, and use the resulting evidence in the final answer. Most existing systems rely on frontier-scale models. We ask whether a compact model can instead learn this behavior through post-training with real tool interaction. We present ProtAgent, a two-stage pipeline that uses supervised fine-tuning (SFT) to learn the tool interface and standard on-policy REINFORCE to improve the interaction policy. Intermediate tool calls receive action rewards, and the final response receives a terminal answer reward. Step-level returns assign these rewards to the responses that produced them and to preceding actions. The episode-level baseline instead assigns one episode return to every action. On a fixed closed-loop protocol with real bioinformatics tools, ProtAgent-8B achieves 61.00% success under the annotation-coverage metric, compared with 41.60% for DeepSeek-V3.2-Thinking 671B and 60.80% for the strongest evaluated frontier system. In the three-seed comparison, mean success is 42.70% with episode-level returns and 61.40% with step-level returns under identity return normalization; the corresponding group-relative results are 36.50% and 58.20%. The metric measures recovery of reference annotations under this protocol; it does not by itself establish biological correctness or complete evidence grounding.