State-Aware Policy Optimization for a Reliable Multi-Turn, Multi-Tool Scientific Agent in Kinetic Biological Models

Motivation

• Known: LLM 기반 agent는 multi-turn reasoning과 tool use를 통해 복잡한 과학적 작업을 수행할 잠재력을 보여왔으며, GRPO 등 RL 기반 방법이 critic 없이도 효율적으로 tool-calling policy를 최적화할 수 있다는 것이 알려져 있다. 또한 SFT나 DPO 등의 방법이 instruction following과 tool-calling 능력 향상에 기여함이 보고되었다.
• Gap: 표준 RL은 stateless 입력을 가정하지만, kinetic biological modeling 같은 과학적 워크플로우는 tool call 간 환경 상태(state)를 유지해야 하는 stateful simulation backend를 필요로 한다. 기존 GRPO 확장 연구들(RC-GRPO, MT-GRPO, RLFactory, AgentFly 등)은 stateless를 가정하거나, state를 재구성하되 per-turn credit assignment를 분리하지 못하거나, static ground truth에 대해 독립적으로 스텝을 평가하는 등 지속적(persistent)이고 stateful한 환경의 문제를 완전히 다루지 못한다.
• Why: 과학 시뮬레이션 도구를 다루는 에이전트는 초기 오류가 후속 단계로 전파(error propagation)되기 쉬워 신뢰성이 매우 중요하며, 이를 해결하면 compact한 모델로도 신뢰할 수 있는 과학적 워크플로우 자동화가 가능해져 QSP 및 생물학적 모델링 연구의 실용성을 크게 높일 수 있다.
• Approach: 본 연구는 multi-turn 상호작용을 per-turn episode로 분해하고, 이전 tool call들을 live simulation backend에서 재실행(replay)하여 환경 상태를 재구성하는 state-aware GRPO 프레임워크를 제안하며, deterministic verification과 frozen LLM judge를 결합한 hybrid reward function으로 검증 가능하고 sample-efficient한 최적화를 수행한다.

Achievement

Figure 2. Pareto frontier of composite score vs. average turn

1. State-aware Per-turn Optimization Framework: SFT와 state-aware GRPO를 결합한 2단계 프레임워크를 제안하여, 환경 상태 재구성과 이전 tool call replay를 통해 multi-turn 상호작용을 per-turn episode로 정식화, stateful 과학적 tool-use 환경에서 세밀한 credit assignment를 가능하게 하고 long-horizon rollout에서의 누적 오류를 줄였다.
2. Optimized Scientific Agent for Biological Models: 제안 프레임워크를 T2B에 적용하여 tool-calling의 일관성과 정확도를 향상시켜 kinetic biological model에 대한 신뢰할 수 있는 질의응답이 가능한 과학 에이전트를 구현했다.
3. Scientific Agent using a Compact Model: Qwen2.5-3B-Instruct 기반의 compact 3B 파라미터 에이전트가 98% tool correctness, 88% argument correctness, 77% task completion을 달성하며, Claude Opus 4.6, GPT-5.4, DeepSeek-V3, O3-mini를 포함한 24개 frontier/open-source 베이스라인을 324-turn 벤치마크(10개 시나리오, 20개 생물학적 모델 기반)에서 능가했다.

How

Figure 1. Overview of the proposed framework, including (a) synthetic benchmark construction and (b,c) staged SFT and st

• Talk2BioModels(T2B)의 데이터 생성 파이프라인을 확장하여 get_modelinfo, simulate_model, steady_state, ask_question, parameter_scan 등 5개 도구를 활용한 multi-turn, multi-tool 시나리오 기반 synthetic benchmark를 구축(self-instruct paradigm 활용, frozen LLM πQA로 질문·답변 생성)
• SFT를 통한 task-specific 모델 초기 미세조정(LoRA adapter 사용)으로 기본적인 tool-calling 능력 확보
• state-aware GRPO를 적용해 multi-turn 상호작용을 per-turn episode로 분해하고, 각 rollout 전에 이전 tool call들을 live simulation backend에서 replay하여 환경 상태를 재구성
• deterministic tool-use verification과 frozen LLM-based judge를 결합한 hybrid reward function(RLVR paradigm)으로 reward-model-free 최적화 수행
• 324-turn, 10개 시나리오, 20개 biological model로 구성된 벤치마크에서 24개 frontier 및 open-source 모델과 tool correctness, argument correctness, task completion 지표로 비교 평가

Originality

• Stateless RL 가정을 넘어, stateful simulation backend를 명시적으로 고려하는 state-aware GRPO를 처음으로 과학적 tool-calling agent 학습에 적용
• 이전 tool call을 live backend에서 replay하여 환경 상태를 재구성함으로써, per-turn credit assignment와 stateful 프로세스의 충실한 실행을 동시에 보존하는 독창적 설계
• deterministic verification과 frozen LLM judge를 결합한 hybrid reward를 통해 reward-model-free이면서도 검증 가능한 RLVR 접근을 QSP/kinetic biological modeling 도메인에 특화하여 적용
• SFT와 state-aware GRPO를 결합한 2단계 파이프라인으로 compact 3B 모델이 frontier 모델들을 능가하는 실용적 성과를 입증

Limitation & Further Study

• 벤치마크가 T2B의 데이터 생성 파이프라인을 확장한 synthetic 데이터에 기반하고 있어, 실제 연구자와의 상호작용에서 발생하는 노이즈나 예측 불가능한 질의 패턴을 충분히 반영하지 못할 가능성이 있음
• 평가가 5개 도구, 10개 시나리오, 20개 생물학적 모델로 제한되어 있어 더 다양한 QSP 워크플로우나 도구 조합에 대한 일반화 가능성 검증이 필요함
• frozen LLM judge를 reward 신호로 사용함에 따라 judge 모델 자체의 편향이나 오류가 최종 정책 학습에 영향을 미칠 수 있음
• 후속 연구로 더 많은 도구, 더 긴 turn 수, 실제 사용자 데이터를 포함한 벤치마크 확장과 다른 과학 도메인(예: 화학, 생물정보학)으로의 프레임워크 일반화 검증이 필요할 것으로 보임

Evaluation

총평: Stateful 과학 시뮬레이션 환경에서의 RL 최적화라는 실질적 문제를 잘 정의하고, compact 모델로 frontier 모델들을 능가하는 인상적인 결과를 보여준 실용적이고 완성도 높은 연구이다.