rStar2-Agent: Agentic Reasoning Technical Report

본 논문은 rStar2-Agent라는 14B 수학 추론 모델을 소개합니다. 이 모델은 Agentic Reinforcement Learning (RL)을 통해 frontier-level 성능을 달성했으며, 기존의 long CoT(Chain-of-Thought)를 넘어선 고급 인지 행동을 보여줍니다. 예를 들어, Python 코딩 도구를 사용하기 전에 신중하게 사고하고, 코드 실행 피드백을 반영하여 복잡한 문제 해결의 중간 단계를 자율적으로 탐색, 검증 및 개선합니다.

이러한 능력은 Agentic RL을 대규모로 효과적으로 만들 수 있는 세 가지 핵심 혁신을 통해 가능해졌습니다.

효율적인 RL 인프라: 높은 처리량의 실행을 지원하고 높은 rollout 비용을 완화하는 신뢰할 수 있는 Python 코드 환경을 포함합니다. 이는 제한된 GPU 자원(64 MI300X GPUs)으로도 훈련을 가능하게 합니다.

GRPO-RoC 알고리즘: Group Relative Policy Optimization with Resampling on Correct의 약자인 이 Agentic RL 알고리즘은 코딩 도구에서 발생하는 본질적인 환경 노이즈를 해결하여 모델이 코드 환경에서 더 효과적으로 추론할 수 있도록 합니다.

효율적인 Agent 훈련 레시피: 비추론적인 Supervised Fine-Tuning (SFT)으로 시작하여 다단계 RL을 거쳐 진행되며, 최소한의 컴퓨팅 비용으로 고급 인지 능력을 발휘합니다.

결과적으로, rStar2-Agent는 510 RL 스텝이라는 짧은 기간(1주일 내)에 사전 학습된 14B 모델을 SOTA 성능으로 끌어올렸으며, AIME24에서 80.6%, AIME25에서 69.8%의 평균 pass@1 점수를 달성하여 671B 규모의 DeepSeek-R1을 능가합니다. 또한, rStar2-Agent-14B는 수학 외에도 Alignment, Scientific Reasoning, Agentic Tool-Use Tasks 등 다양한 영역으로 강력한 일반화 능력을 보여줍니다.

Agentic Reinforcement Learning Methodology

rStar2-Agent는 LLM이 전용 실행 환경 내에서 Python 코딩 도구를 Native하게 사용하여 복잡한 문제 해결을 위한 고급 및 효과적인 추론을 가능하게 합니다.

2.1 Smarter Reasoning in a Code Environment

모델은 Python 코드 환경에서 인간과 유사한 인지 행동을 보여주도록 훈련됩니다: (i) 적절한 추론 단계에서 도구를 호출하고, (ii) 논리적으로 정확하고 기능적인 코드를 작성하며, (iii) 실행 결과를 신중하게 반영하여 후속 추론 단계를 안내합니다.

* Multi-turn Rollout: 모델은 코드 환경과의 여러 인터랙티브 턴을 통해 전체 trajectory를 생성합니다. 첫 번째 턴은 시스템 프롬프트와 질문으로 시작되며, 모델은 assistant 역할로 초기 추론 trajectory를 생성하고 EOS 토큰으로 끝냅니다. 만약 코드 툴 호출이 있다면, 해당 코드 블록은 추출되어 환경 서비스에 의해 실행됩니다. 실행 결과는 user 역할 아래 $<tool_response>$ 태그 내에 trajectory에 추가되며, 모델은 이 업데이트된 context를 입력으로 받아 다음 턴의 추론을 계속합니다. 이 Multi-turn rollout 프로세스는 모델이 최종 답변을 생성하거나 미리 정의된 최대 턴 수 T에 도달할 때까지 반복됩니다.
* Tool Call Format: 도구 호출을 위해 구조화된 JSON 형식인 $<tool_call>{"name": "...", "arguments": {"code": "..."}}</tool_call>$을 사용합니다. 환경 서비스는 성공적인 실행(출력 포함/미포함), 실행 오류(에러 메시지 및 traceback 로그), 또는 타임아웃(구문은 유효하나 시간 초과)의 네 가지 가능한 응답 중 하나를 반환합니다. 모든 경우에 환경 피드백은 $<tool_response>$ 태그에 래핑되어 모델에 피드백됩니다.
* Prompt Template: 모델은 $<reason>...</reason>$ 내에 추론 과정을, $<answer>...</answer>$ 내에 최종 답변을 생성하도록 지시받습니다. 최종 결과는 $\boxed{}$ 안에 래핑되어야 합니다. Multi-turn rollout 설정에서는 여러 $<reason>$ 블록이 생성될 수 있지만, $<answer>$ 블록은 하나만 허용됩니다.

2.2 End-to-End Agentic Reinforcement Learning

2.2.1 Preliminary: GRPO

Group Relative Policy Optimization (GRPO)는 이전 정책 $\pi_{\theta_{old}}$에서 롤아웃 trajectories 그룹 $\{o_1, o_2, \dots, o_G\}$를 샘플링하고 다음 목적 함수를 최대화하여 정책 $\pi_\theta$를 최적화합니다:
$$J_{GRPO}(\theta) = E_{(q,a)\sim D, \{o_i\}_{i=1}^G (\cdot|q) \sim \pi_{\theta_{old}}} \left[ \frac{1}{G} \sum_{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \min\left( \frac{\pi_\theta(o_{i,t}|q, o_{i,<t})}{\pi_{\theta_{old}}(o_{i,t}|q, o_{i,<t})} A_{i,t}, \text{clip}\left( \frac{\pi_\theta(o_{i,t}|q, o_{i,<t})}{\pi_{\theta_{old}}(o_{i,t}|q, o_{i,<t})}, 1-\epsilon, 1+\epsilon \right) A_{i,t} \right) - \beta D_{KL}(\pi_\theta \parallel \pi_{ref}) \right]$$
여기서 $A_{i,t}$는 다음과 같이 계산되는 Advantage입니다:
$$A_{i,t} = \frac{r_i - \text{mean}(\{r_1, r_2, \dots, r_G\})}{\text{std}(\{r_1, r_2, \dots, r_G\})}$$
$r_i$는 롤아웃 trajectory $o_i$에 할당된 이진 정확도 보상입니다:
$$r_i = \begin{cases} 1 & \text{if is equivalent}(a, o_i) \\ 0 & \text{otherwise} \end{cases}$$
본 연구에서는 KL divergence penalty를 제거하고, importance sampling ratio의 상한을 완화하며 (Clip-Higher 전략, $\epsilon_{high}$를 0.2에서 0.28로 증가), 엔트로피 손실 항을 제거하여 더 많은 탐색을 가능하게 합니다.

2.2.2 Challenges in Agentic Reinforcement Learning

코딩 도구와 코드 환경은 추론 과정에 내재된 노이즈를 도입합니다. 모델이 구문적 또는 논리적으로 잘못된 코드를 생성하면 환경에서 반환되는 오류 메시지(예: error message)가 모델을 오도하여 추론을 진행하기보다는 오류를 수정하는 데 귀중한 토큰을 낭비하게 할 수 있습니다. Outcome-only reward 방식에서는 최종 답변이 정확하다면 중간 단계의 오류가 있는 trajectory도 양의 보상을 받게 되어 모델이 이러한 오류를 허용 가능한 것으로 간주하게 됩니다.

2.2.3 GRPO-RoC: Group Relative Policy Optimization with Resampling on Correct

더 효과적인 Agentic RL을 위해 GRPO-RoC를 제안합니다. 이는 Resample on Correct (RoC) rollout 전략을 통해 낮은 품질의 노이즈가 많은 trajectory를 효과적으로 필터링합니다.

* Design Principle: Answer-only Outcome Reward: 보상 해킹을 피하기 위해 Eq. 3과 같이 최소한의 answer-only outcome reward를 사용합니다.
* Resample on Correct (RoC): 이 rollout 전략은 먼저 2G개의 롤아웃 trajectory를 Oversample한 다음, 표준 배치 크기인 G개로 Downsample합니다. 긍정적인 trajectory는 도구로 인한 오류나 형식 문제(tool call errors, format violations)가 최소화된 최고 품질의 것들만 필터링하여 유지합니다. 부정적인 trajectory는 균등하게 Downsample됩니다. 이 비대칭 샘플링은 다양한 실패 모드를 유익한 음성 신호로 유지하면서 더 높은 품질의 성공 사례에 중점을 두어 긍정적인 지도 학습을 강화합니다.
* Negative samples: Oneg에서 $\lfloor |Oneg|/2 \rfloor$개의 trajectory를 샘플링하여 실패의 다양성을 유지합니다.
* Positive samples: Opos에서 절반의 trajectory를 샘플링하며, 더 높은 품질의 trace에 우선순위를 둡니다. 각 trajectory는 tool call errors ($p_{err}$) 및 format violations ($p_{format}$) 두 가지 유형의 중간 문제에 대해 점수가 매겨집니다.
* $p_{err} = \frac{\text{num of error tool calls}}{\text{num of all tool calls}}$ (도구 호출이 없으면 0.5)
* $p_{format} = \begin{cases} 1 & \text{if no <answer> tags} \\ \min(1, \frac{\text{num of <answer> tags}-1}{\text{num of turns}}) & \text{otherwise} \end{cases}$
* 총 페널티 점수 $p_{total} = p_{err} + p_{format}$를 계산하고, $p_{total}$에 반비례하는 확률로 Positive trajectory를 샘플링합니다.

GRPO-RoC의 최종 RL 목적 함수는 다음과 같이 공식화됩니다:
$$J_{GRPO-RoC}(\theta) = E_{(q,a)\sim D, \{o_i\}_{i=1}^{2G} (\cdot|q) \sim \pi_{\theta_{old}}} \left[ \frac{1}{G} \sum_{\hat{o}_i \in \{\hat{o}_i\}_{i=1}^G} \frac{1}{|\hat{o}_i|} \sum_{t=1}^{|\hat{o}_i|} \min\left( \frac{\pi_\theta(\hat{o}_{i,t}|q, \hat{o}_{i,<t})}{\pi_{\theta_{old}}(\hat{o}_{i,t}|q, \hat{o}_{i,<t})} A_{i,t}, \text{clip}\left( \frac{\pi_\theta(\hat{o}_{i,t}|q, \hat{o}_{i,<t})}{\pi_{\theta_{old}}(\hat{o}_{i,t}|q, \hat{o}_{i,<t})}, 1-\epsilon, 1+\epsilon \right) A_{i,t} \right) \right]$$
여기서 $\{\hat{o}_i\}_{i=1}^G$는 RoC 전략에 따라 선택된 G개의 trajectory를 나타냅니다.

3 Large-Scale Agentic RL Infrastructure

* 3.1 Reliable High-Throughput Code Environment: 45K 동시 툴 호출을 처리할 수 있는 고처리량, 격리된 코드 환경을 구축하여 평균 0.3초 만에 실행 피드백을 반환합니다.
* 3.2 Load-Balanced Rollout Scheduler: RL rollout의 비효율성을 해결하기 위해 사용 가능한 KV cache 용량에 따라 rollout 요청을 동적으로 할당하는 Load-Balanced Rollout Scheduler를 도입하여 계산 활용도를 최대화합니다. 이 인프라 덕분에 64개의 MI300X GPU로 1주일 만에 훈련을 완료했습니다.

4 Training Recipe

* 4.1 Non-Reasoning Cold Start for Instruction Following: 추론 관련 SFT 대신 일반적인 Instruction Following, 코딩 도구 사용 및 포매팅을 주입하는 SFT 단계로 시작합니다. 이는 SFT 과적합을 피하고 초기 응답 길이를 짧게 유지하여 RL이 추론 능력을 더 효과적으로 배양할 수 있도록 합니다.
* 4.2 RL Data Curation: 훈련 데이터를 Math Problem Solving, Tool Use 및 General Instruction Following 세 가지 범주로 분류합니다. Math Problem Solving 데이터셋은 DeepSeek-Math에서 파생되었으며 AIME/AMC/MATH와 같은 고급 수학 문제로 구성됩니다. Tool Use 데이터셋은 GPT-4를 사용하여 생성된 Python 코드와 이에 상응하는 실행 결과를 포함하는 instruction-following 예제입니다. General Instruction Following 데이터셋은 광범위한 공개 소스 SFT 데이터셋의 하위 집합입니다.
* 4.3 Multi-Stage RL Training: GRPO-RoC와 함께 다단계 RL 훈련을 수행하며, 점진적으로 작업 난이도와 최대 훈련 길이를 늘립니다. 각 단계를 짧은 길이(8K→12K)로 제한하여 RL 비용을 크게 줄이고 효율적인 추론 전략을 장려합니다. 단 510 RL 스텝 만에 모델은 frontier-level 수학 추론을 빠르게 달성하여 높은 능력과 탁월한 훈련 효율성을 입증했습니다. 초기 시도에서 Multi-stage RL의 실패 경험을 통해, 각 RL 스테이지에서 max_seq_len을 점진적으로 늘리고, RL 훈련 시작 시 max_turn을 1로 제한하며, Reward Model을 사용하지 않는 것이 중요함을 강조합니다. 또한, 코드 인터프리터 오류 메시지나 타임아웃 오류 메시지에 대해 명시적인 Reflection Token을 학습하는 것이 모델의 인지적 행동 발달에 필수적임을 발견했습니다.

5 Experiments

* 5.1 Setup: 모델은 Qwen2.5-14B-Chat을 기반으로 하며, 총 훈련 토큰은 2.8T입니다. 훈련은 64 MI300X GPU에서 수행되었으며, 훈련 시간은 총 1주일이었습니다.
* 5.2 rStar2-Agent-14B Main Results: AIME24에서 80.6%의 정확도를 달성하여 OpenAI o3-mini (medium), DeepSeek-R1, Claude Opus 4.0 (thinking)을 능가했습니다. AIME25에서 69.8%, HMMT25에서 52.7%를 기록했습니다. 이는 Qwen3-14B-Chat을 기반으로 하는 14B 모델임에도 불구하고, 671B DeepSeek-R1과 같은 훨씬 큰 모델을 능가하는 인상적인 성능을 보여줍니다.
* 5.3 Ablation Study and Discussions:
* GRPO-RoC의 중요성: GRPO-RoC 없이 훈련된 모델은 성능 저하와 함께 Tool Call 오류 비율이 더 높았습니다. 이는 RoC가 모델이 Tool-Use 오류를 줄이도록 효과적으로 학습시키는 데 중요함을 시사합니다.
* Multi-Stage RL의 이점: Multi-Stage RL 훈련이 단일 단계 훈련보다 더 효과적이라는 것을 보여주었으며, 점진적인 max_seq_len 증가가 핵심입니다.
* Reflection Token의 영향: Reflection Token을 도입하여 모델이 환경 피드백에 대해 명시적으로 성찰하도록 장려한 결과, 문제 해결의 성공률과 효율성이 향상되었습니다.
* 5.4 Analysis of Agentic Reasoning Behaviors: rStar2-Agent는 단순한 긴 CoT를 넘어선 고급 인지 행동을 보입니다. 모델은 Tool-Use 오류 발생 시 이를 인식하고, 문제를 해결하기 위해 코드를 수정하거나, 오류로부터 교훈을 얻어 새로운 접근 방식을 시도하는 등의 Self-correction 및 Self-reflection 능력을 보여줍니다.

6 Conclusion

rStar2-Agent는 Agentic RL을 통해 14B 모델을 frontier-level 수학 추론 능력으로 끌어올렸으며, 이는 효율적인 RL 인프라, GRPO-RoC 알고리즘, 그리고 효율적인 다단계 훈련 레시피 덕분입니다. 특히, GRPO-RoC는 코딩 환경의 노이즈 문제를 해결하고 모델이 고품질의 추론 경로를 학습하도록 돕습니다. 모델은 수학뿐만 아니라 Alignment, Scientific Reasoning, Tool-Use 등 다양한 분야에서 강력한 일반화 능력을 보여줍니다.