rStar2-Agent: 探索代理推理在数学问题解决中的应用

大家好,我是这篇博客的作者。今天,我想和大家聊聊 rStar2-Agent 这个项目。它是一个专注于代理强化学习(agentic reinforcement learning)的数学推理模型,能帮助模型在解决复杂问题时更聪明地思考,而不是单纯地延长思考链条。如果你对大型语言模型(LLM)如何提升数学推理能力感兴趣,尤其是如何通过工具调用来验证和优化步骤,这篇文章可能会给你一些启发。我们一步步来探讨,从基础介绍到实际操作,再到背后的技术细节。

想象一下,你在解决一个棘手的数学问题,比如一个高中级别的竞赛题。你可能会先试着一步步推理,但如果中间卡住了,怎么办?rStar2-Agent 就是设计来处理这种场景的。它让模型学会使用 Python 编码工具来探索想法、验证计算,并从执行反馈中学习。这不是简单的链式思考(Chain-of-Thought),而是一种更灵活的代理式方法,能让一个 14B 参数的模型达到甚至超过一些更大模型的表现。

rStar2-Agent 是什么?

rStar2-Agent 是一个 14B 的数学推理模型,通过代理强化学习训练而成。它强调“聪明思考”而非“长思考”。简单说,这个模型能规划推理、使用编码工具来验证想法,并在过程中反思和调整。背后的关键创新包括一个高效的强化学习算法 GRPO-RoC(Group Relative Policy Optimization with Resample-on-Correct),一个支持高吞吐量工具执行的 RL 基础设施,以及一个从非推理监督微调(SFT)开始的多阶段 RL 训练配方。

为什么这个重要?因为在数学问题上,它能在有限的计算资源下达到前沿水平。比如,在 AIME24 上平均通过率 80.6%,在 AIME25 上 69.8%,甚至超过了 671B 的 DeepSeek-R1,而且响应更短。除了数学,它在对齐、科学推理和代理工具使用任务上也表现出色。

如果你好奇“代理强化学习”具体怎么工作,它本质上是让模型像代理一样与环境互动。这里环境是 Python 解释器,模型生成代码、执行、获取反馈,然后优化自己的策略。这能帮助模型处理复杂问题,避免长 CoT 中的常见错误。

Figure 1: rStar2-Agent-14B reaches frontier-level math reasoning in just 510 RL training step

如上图所示,这个模型只需 510 个 RL 训练步就能达到高水平表现。

rStar2-Agent 的最新进展

项目有一些值得注意的更新:

  • 2025 年 7 月 15 日:发布了 rStar-Coder 论文和数据集,包括 418K 个竞赛级代码问题,带测试用例,帮助小模型提升代码推理。
  • 2025 年 2 月 10 日:正在招聘实习生,如果你对改善 LLM 推理感兴趣,可以联系相关邮箱。
  • 2025 年 1 月 21 日:开源了 rStar-Math 代码。
  • 2025 年 1 月 9 日:发布了 rStar-Math 论文。

此外,先前工作如 Mutual Reasoning 和 rStar-Math 有独立的开源分支。

这些更新显示项目在持续演进,从数学到代码,再到更广的推理领域。

如何上手 rStar2-Agent 并使用工具调用?

如果你想亲手试试 rStar2-Agent,尤其是带工具调用的交互,这里有详细步骤。我们从安装开始,一步步走。

安装指南

有两种安装方式,选择适合你的。

方式 1: 手动安装

  1. 初始化并更新子模块:

    git submodule init
    git submodule update
    
  2. 安装 verl:

    pip install "torch<2.8"
    pip install -r verl/requirements_sglang.txt
    pip install -e verl
    
  3. 安装 code judge:

    pip install -r code-judge/requirements.txt
    pip install -e code-judge
    
  4. 安装 rstar2_agent:

    pip install -e .
    

方式 2: 自动安装

运行脚本:

bash install.sh

安装后,确保环境安全,尤其是 code judge 会执行任意代码,最好用 Docker 隔离。

设置 Code Judge 服务器

Code Judge 是 rStar2-Agent 的工具调用服务器,用于执行生成的 Python 代码。注意:部署在隔离环境中,避免暴露到外部网络。

  1. 启动 Redis 服务器:

    redis-server --daemonize yes --protected-mode no --bind 0.0.0.0
    
  2. 启动主服务器(仅在主节点):
    用 tmux 运行:

    tmux new-session -d -s server \
      'cd $WORKSPACE/code-judge && \
       MAX_EXECUTION_TIME=4 \
       REDIS_URI="redis://$MASTER_ADDR:6379" \
       RUN_WORKERS=0 \
       uvicorn app.main:app --host 0.0.0.0 --port 8088 --workers 16 \
       2>&1 | tee server.log'
    

    替换 MASTER_ADDR 为实际路径。

  3. 启动 workers(可在多节点上):

    tmux new-session -d -s worker \
      'cd $WORKSPACE/code-judge && \
       MAX_EXECUTION_TIME=4 \
       REDIS_URI="redis://$MASTER_ADDR:6379" \
       MAX_WORKERS=64 \
       python run_workers.py \
       2>&1 | tee worker.log'
    

    根据 CPU 调整 MAX_WORKERS。

启动 VLLM 服务器

VLLM 用于服务模型:

vllm serve /path/to/your/model \
    --host 0.0.0.0 \
    --port 8000 \
    --enable-auto-tool-choice \
    --tool-call-parser hermes

验证服务器:

curl http://localhost:8000/v1/models

运行互动聊天带工具调用

用脚本测试:

python examples/chat_with_tool_call.py \
    --model /path/to/your/model \
    --prompt "Solve the system of equations: 2x + 3y = 7, x - y = 1" \
    --max_tokens 8192

选项包括模型路径、提示和最大 token。

通过这些步骤,你就能看到模型如何使用工具解决方程组,比如生成代码验证解。

如何评估 rStar2-Agent 的性能?

评估是验证模型的关键。这里用数学基准如 AIME 2024/2025 和 MATH500。

环境设置

参考上面的安装和 Code Judge 设置。

运行评估脚本

对于 AIME:

MODEL_PATH=/path/to/your/model bash examples/aime_eval.sh

对于 MATH500:

MODEL_PATH=/path/to/your/model bash examples/math500_eval.sh

这些脚本会输出通过率,帮助你比较模型表现。

rStar2-Agent 的 RL 训练过程是怎样的?

现在,我们聊聊训练部分。如果你想自己训练一个类似模型,这里有完整指南。rStar2-Agent 的 RL 框架基于 Verl 和 Code Judge,支持从 SFT 后的模型开始训练。

环境设置

同安装和 Code Judge 部分。

数据准备

用 DAPO-17k(英文子集)作为训练数据,AIME24 作为测试。

  1. 处理 AIME 2024:

    python data_preprocess/aime2024_rstar2_agent_loop.py
    
  2. 处理 DAPO:

    python data_preprocess/dapo_rstar2_agent_loop.py
    

模型设置

下载基模型 Qwen3-14B-Base:

huggingface-cli download Qwen/Qwen3-14B-Base --local-dir $HOME/models/Qwen3-14B-Base

注意:基模型需先做指令跟随 SFT 以优化 RL 性能。

运行训练

对于 8x A100/H100 GPUs:

bash examples/run_qwen3-14b_rstar2_agent_weave.sh

调整硬件参数。

配置数据增强

框架支持采样策略提升效率:

  • 启用下采样:augmentation.do_down_sampling=True
  • 目标轨迹数:augmentation.down_sampling_config.down_sample_to_n=16
  • 拒绝等奖励采样:augmentation.down_sampling_config.reject_equal_reward=True
  • 按工具调用错误率重采样正确轨迹:augmentation.down_sampling_config.roc_error_ratio=True
  • 按答案格式重采样:augmentation.down_sampling_config.roc_answer_format=True
  • 最小负轨迹:augmentation.down_sampling_config.min_zero_reward_trace_num=2
  • 最小正轨迹:augmentation.down_sampling_config.min_non_zero_reward_trace_num=2

这些设置帮助模型从高质量轨迹中学习。

重要提示:当前框架迁移到 VERL v0.5 以兼容社区标准。原始框架有更多特性,如负载均衡调度,但这个版本已验证前 50 步训练相似。我们计划未来迁移所有功能。如果遇到问题,欢迎反馈。

常见问题排查

  • Redis 连接错误:检查 Redis 是否运行并可访问。
  • GPU 内存问题:调整批次大小。
  • Code Judge 超时:增加 MAX_EXECUTION_TIME。
  • Worker 扩展:根据 CPU 核调整 MAX_WORKERS。

日志在 code-judge 的 server.log 和 worker.log,以及训练输出目录。

rStar2-Agent 的核心方法:代理强化学习

让我们深入谈谈技术。rStar2-Agent 通过代理 RL 让模型在代码环境中更聪明地推理。

代码环境中更聪明推理

模型不只是长 CoT,它会仔细思考前使用 Python 工具,反思执行反馈,自主探索、验证和精炼步骤。这处理长 CoT 无法解决的微妙错误或需要创意转变的问题。

端到端代理 RL

基于 GRPO,但针对代理挑战优化。

GRPO 初步

GRPO 是相对策略优化,使用组内相对优势。

代理 RL 挑战

代码工具引入噪声,如语法错误导致浪费 token。现有 RL 只看最终奖励,忽略中间失败。

GRPO-RoC:组相对策略优化带正确重采样

RoC 先超采样轨迹,然后下采样。正轨迹保留高质量(最小错误、格式问题)的,负轨迹均匀下采样。这稳定训练,避免奖励黑客风险,让模型从干净正例中学习更好工具使用和认知能力。

大规模代理 RL 基础设施

支持高效训练。

可靠高吞吐代码环境

处理 45K 并发工具调用,平均 0.3 秒反馈。

负载均衡 rollout 调度器

动态分配请求基于 KV 缓存容量,最大化 GPU 利用。

训练配方

从预训练 14B 模型到前沿水平,只需 510 RL 步,一周内完成。

非推理冷启动指令跟随

从非推理 SFT 开始,确保基本指令能力。

RL 数据 curation

用 DAPO 等数据集,逐步增加难度。

多阶段 RL 训练

从短响应开始,多阶段 RL,逐步复杂。

不成功尝试和教训

早期实验显示直接从长 CoT 开始不佳;需渐进。

实验结果

设置

用 64 MI300X GPUs 训练。

rStar2-Agent-14B 主结果

模型 AIME24 AIME25 HMMT25
OpenAI o3-mini (medium) 79.6 77.0 53.0
DeepSeek-R1 (671B) 79.8 70.0 44.4
DeepSeek-R1-Zero (671B) 71.0 53.3 46.0
Claude-Opus-4.0 (Think) 76.0 69.2
QWQ-32B 79.5 65.8 47.5
rStar2-Agent-14B 80.6 69.8 52.7

超越更大模型,响应更短。泛化到对齐、科学推理、代理工具使用。

消融研究和讨论

GRPO-RoC 关键;移除 RoC 性能下降。

代理推理行为分析

模型展示高级认知:高熵 token 用于反思、探索错误。工具调用 token 低熵,因预训练代码语料。

结语

rStar2-Agent 展示了代理 RL 在提升小模型推理上的潜力。通过高效基础设施、GRPO-RoC 和多阶段训练,它在数学上达到前沿,并泛化更广。代码和配方开源,欢迎探索。

FAQ

rStar2-Agent 如何比长 CoT 更好?

长 CoT 易中间错误;rStar2-Agent 用工具验证,反思反馈,更灵活。

训练 rStar2-Agent 需要多少资源?

64 MI300X GPUs,一周内 510 步。

如何处理代码执行错误?

模型生成高熵 token 诊断、探索替代,精炼推理。

rStar2-Agent 只用于数学吗?

不,它泛化到对齐、科学推理、代理工具使用。

为什么用 GRPO-RoC?

它处理代码噪声,强调高质量正轨迹,提高稳定性。

如何下载模型?

从 Hugging Face 下载 Qwen3-14B-Base,然后训练。

评估基准是什么?

AIME24/25、HMMT25、MATH500。

安装遇到问题怎么办?

检查 Redis、调整批次、查看日志。

rStar2-Agent 的创新点?

GRPO-RoC、高吞吐环境、多阶段 RL 配方。