代理式 AI/生成式 AI

如何使用强化学习训练科学智能体

科学过程可能重复且繁琐,研究人员往往需要花费数小时深入研读论文、管理实验流程或整理庞大的多模态数据集。科学 AI 智能体能够承担大量繁重任务,作为协助审查文献、生成假设、规划实验、提交计算作业、编排实验室操作、分析结果并总结发现的助手。这使得研究人员得以专注于创造性思考与科学探索。

但构建科学 AI 助手极具挑战性。智能体需在多个研究步骤中制定高层计划,同时整合内存与上下文管理。一个错误便可能影响整个研究任务。此外,通用大语言模型难以有效运用特定领域的工具,尤其在前沿研究领域更为明显。通过计算或现实世界数据验证结果往往耗时较长,要求智能体在数小时、数天甚至更长时间内保持行为一致。

NVIDIA NeMo Gym 和 NeMo RL 作为 NVIDIA NeMo 框架 套件中的开源库,提供了一个统一且模块化的强化学习堆栈,可用于在包括科学研究在内的各个领域构建可靠的代理式 AI。NeMo Gym 使开发者能够创建逼真的环境,智能体可在其中进行交互、学习并完成特定领域的任务,从而生成高质量、可验证且面向特定领域的推出数据。这些训练数据随后可与 NeMo RL 结合使用,以高效地大规模适应和优化智能体性能,NeMo RL

这两个库在最新 Nemotron-3-Nano 的后训练过程中发挥了关键作用,Nemotron-3-Nano 是一种针对目标任务优化的经济高效模型,能够在低推理成本下提供高精度。

Edison Scientific 是一位使用 NeMo Gym 和 NeMo RL 的开发者,该公司致力于推动科学发现的自动化。作为非营利研究组织 FutureHouse 的衍生企业,该公司利用该基础设施为 Aviary 提供支持,Aviary 是一个涵盖生物学、化学及相关领域的科学强化学习训练环境框架。

在本博客中,我们将演示如何利用 NeMo Gym 构建代理式训练环境,并结合 NeMo RL 进行训练。以 Aviary 为例,展示其在特定领域科学强化学习环境中的应用。

强化学习如何拓展科学领域中大语言模型的功能

并非所有大语言模型(LLM)都能执行复杂的科学工作流程。预训练使模型学会预测下一个 token,从而构建起广泛的知识基础,但并不具备特定领域的专业技能。这一基础使得模型能够在结构化事实类问题(例如基因与疾病关联、药物作用机制或临床时间线)上实现零样本表现。随后的后训练阶段则教会模型遵循指令,并通过迭代优化与对齐过程,使其输出更符合特定领域的偏好与要求。

后训练通常从监督式微调(SFT)开始,模型通过指令-响应对,采用下一词预测的对数概率损失进行学习。该过程高度依赖高质量的专家数据或经过筛选的合成数据,且对数据中的错误较为敏感。SFT 的性能受限于其训练数据集的覆盖范围,即使存在其他正确的输出(例如不同的有效代码实现),损失函数也仅对复现参考答案给予奖励。

训练流程有时会引入强化学习(RL),以拓展模型的推理与行动能力,使其超越监督数据的局限。RL 在训练过程中利用奖励函数对模型或策略的输出进行评分。在基于人类反馈的强化学习(RLHF)中,人类根据自身偏好或评级对模型的回答进行排序。基于 AI 反馈的强化学习(RLAIF)则使用大语言模型(LLM)充当评判者,从而省去人工偏好标注的环节。而带有可验证奖励的强化学习(RLVR)通过代码执行等计算性验证方式,生成客观且可复现的奖励信号。

RLVR 有助于训练科学智能体,因为它使模型能够设计并运行实验、评估结果,并通过验证设计与奖励塑造来优化科学指标。科学强化学习可在多步环境中运行,在该环境中,智能体采取行动、观察反馈,并持续迭代,直至任务完成。训练过程可基于完整轨迹或单个状态转换进行。借助强化学习,科学智能体能够整合在预训练和监督微调阶段所掌握的技能,构建新的工作流程,以实现特定的科学目标。

NeMo Gym 与 NeMo RL 如何优化代理式训练与评估

为 LLM 智能体实现 RL 需要一个训练框架和环境,用于定义智能体可执行的动作、可观察的状态以及通过行动获得的奖励。训练框架(例如 NeMo RL)负责运行训练算法(例如组相对策略优化(GRPO)),管理模型部署与验证所需的计算资源,并协调模型权重的更新。新版 NeMo RL 支持 on-policy 蒸馏、asyncRL、高级 RL 算法以及端到端 FP8 RL 训练。

智能体通过采取行动并调用必要工具,与环境进行交互循环;环境则提供行动的观察结果与奖励,维持状态的持久性,并判断任务是否完成。环境可包含简单的 Python 执行沙盒,也可涵盖用于评估分子克隆等工作流程的完整研究软件堆栈。

训练 AI 科学家需要在诸多复杂任务中表现出色的模型。在实践中,这意味着模型需在文献合成、假设生成、实验设计和数据分析等应用场景中完成成百上千项不同任务,而每一类应用都对应着独特的验证逻辑。随着任务多样性的增加,由于各类依赖关系及特定领域的要求,训练环境的基础设施管理也变得更加复杂。为应对这一挑战,我们开发了 NeMo Gym,这是一个用于大规模构建强化学习训练环境的开源框架。

NeMo Gym 是 LLM 后训练中用于 RL 数据、环境和奖励信号的核心平台。它提供基础设施,支持训练环境的开发、部署集合的扩展,并可与您首选的训练框架无缝集成。环境相互隔离并通过 REST API 暴露,从而实现并行执行和可扩展部署,同时避免依赖冲突。

NeMo Gym 提供三种核心服务器抽象。训练环境通常包含协同工作的全部三种服务器类型:

  • 模型: 包含兼容 OpenAI 的端点,支持推理与工具调用功能。模型可于本地或云端运行,并能与多种后端(包括 OpenAI、Azure 和 vLLM)协同工作。该抽象层将模型部署与智能体逻辑相分离。
  • 资源:提供可通过工具调用及验证逻辑(用于评估任务表现)访问的工具实现。此抽象机制可将繁重的处理任务卸载,使智能体能够异步调用模型进行推理,同时调用资源执行工具操作与验证。
  • 智能体:负责编排模型与资源间的交互,实现请求路由、协调多轮对话,并统一响应格式。

NeMo Gym 可从复杂的训练环境中生成结果与奖励,从而形成强化学习训练所需的优化目标。NeMo Gym 能够与现有的环境、系统及强化学习训练框架互操作,支持用户使用自定义环境或 NVIDIA 精心设计的环境进行大语言模型的后训练。当与用于训练算法和基础设施的 NeMo RL 结合使用时,这两个库可为代理式训练与强化学习提供可扩展的工作流程。

NeMo Gym 的实际应用:在 Edison Scientific 训练科学推理智能体

Edison Scientific 正在利用 NeMo Gym 和 NeMo RL 扩展 AI 智能体,以实现科学发现的自动化。其中包括 Aviary,该平台能够训练生物学、化学及相关领域的智能体,可执行 文献研究生物信息数据分析实验室任务解决分子克隆问题,以及处理多步骤的科学问题。

Aviary 负责管理 RL 环境的状态、工具执行、奖励以及观察格式。其开源库包含适用于数学、科学文献研究和数据分析的环境。NeMo Gym 在 Aviary 上运行,由 Aviary 控制环境逻辑,同时 NeMo Gym 提供可扩展的部署集合、NVIDIA 精心策划的其他训练环境,并与 NeMo RL 集成,支持大规模训练。

每个 Aviary 环境都实现了两种核心方法:reset() 和 step()。reset() 方法用于初始化环境,返回初始观察结果以及可用的工具列表。step() 方法执行动作,返回新的观察结果、奖励,以及终止或截断信号。动作表现为工具请求,可能包含多个工具调用。

Edison Scientific 正在利用 NeMo Gym 结合 Aviary 训练用于 Jupyter Notebook 数据分析的智能体,以执行生物信息学任务。在每个步骤中,智能体会查看 Notebook 并编辑其中的单元格。由于 Notebook 的规模可能超出模型的上下文窗口,Edison Scientific 引入了两项机制来控制上下文的增长:一是清除交互历史,使智能体仅能看到原始指令、此前的所有操作以及当前的 Notebook 内容;二是调整 GRPO 的分组方式,使其在单个步骤上进行操作,而非在整个轨迹上操作。这一改进支持在状态转换过程中进行训练,有效缩短了上下文长度,并实现了分步的奖励反馈机制。
作为验证平台,Edison Scientific 在 Aviary 中构建了一个基于 Jupyter 的数据分析环境,将其与 NeMo Gym 集成,并开发了一个名为 BixBench 的可验证生物信息学问题基准测试。

Software architecture diagram showing the relationship between NeMo Gym, NeMo RL, and downstream scientific tasks. NeMo Gym provides agents, resources, and model interfaces that generate training signals from scientific and synthetic data. NeMo RL supplies the models, RL algorithms, and training toolkit used to optimize agent behavior. Outputs are shown flowing to downstream evaluation and inspection, though this represents a conceptual sketch rather than a fully deployed end-to-end system. Software architecture diagram showing the relationship between NeMo Gym, NeMo RL, and downstream scientific tasks. NeMo Gym provides agents, resources, and model interfaces that generate training signals from scientific and synthetic data. NeMo RL supplies the models, RL algorithms, and training toolkit used to optimize agent behavior. Outputs are shown flowing to downstream evaluation and inspection, though this represents a conceptual sketch rather than a fully deployed end-to-end system.
图 1。NeMo Gym 与 NeMo RL 的高层架构,用于训练面向科学任务的强化学习智能体。该示意图展示了由 NeMo Gym 管理的训练环境如何与 NeMo RL 训练基础设施相连接,以训练可支持下游科学工作流的智能体。

在 NeMo Gym 中构建用于训练或下游任务的代理式环境非常简单,只需几个步骤。

第 1 步:安装 NeMo Gym

克隆 NeMo Gym 仓库,安装 uv Python 包,然后创建虚拟环境:

# Clone the repository
git clone git@github.com:NVIDIA-NeMo/Gym.git
cd Gym

# Install UV (Python package manager)
curl -LsSf https://astral.sh/uv/install.sh | sh
source $HOME/.local/bin/env

# Create virtual environment
uv venv --python 3.12
source .venv/bin/activate

# Install NeMo Gym
uv sync --extra dev --group docs

第 2 步:配置模型

您可以使用 OpenAI 等托管模型,也可以通过 NVIDIA NIM 或 vLLM 在本地部署模型。在本示例中,我们将使用 HuggingFace 上的 nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 模型,并通过支持工具调用的 vLLM 进行部署。有关如何在 vLLM 中使用该模型的更多详细信息,请参阅此指南

pip install -U "vllm>=0.12.0"

wget https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/resolve/main/nano_v3_reasoning_parser.py

vllm serve nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
 --max-num-seqs 8 \
  --tensor-parallel-size 1 \
  --max-model-len 262144 \
  --port 10240 \
  --trust-remote-code \
  --tool-call-parser qwen3_coder \
  --enable-auto-tool-choice \
  --reasoning-parser-plugin nano_v3_reasoning_parser.py \
  --reasoning-parser nano_v3

然后,在 NeMo Gym 根目录下创建 env.yaml 文件:

policy_base_url: http://localhost:10240/v1
policy_api_key: EMPTY
policy_model_name: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16

第 3 步:在 NeMo Gym 中使用基础智能体测试 Aviary 环境

现在,让我们在 GSM8K 环境中运行智能体,这是一个包含数学问题的数据集,智能体可在其中使用计算器工具。

在 NeMo Gym 中,`ng_run` 命令用于启动服务器。为配置服务器,需提供相应的配置文件。此处提供了两个配置文件:`gsm8k_aviary.yaml` 用于配置资源服务器与代理服务器,`vllm_model.yaml` 用于定义模型服务器。

ng_run "+config_paths=[resources_servers/aviary/configs/gsm8k_aviary.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml]"

所有服务器启动后,您应收到类似于以下内容的日志:

All 3 / 3 servers ready! Polling every 60s

####################################################################################################
#
# Server Instances
#
####################################################################################################

[1] gsm8k_aviary_resources_server (resources_servers/aviary)
{
    'process_name': 'gsm8k_aviary_resources_server',
    'server_type': 'resources_servers',
    'name': 'aviary',
    'dir_path': (
        '/home/ubuntu/Gym/resources_servers/aviary'
    ),
    'entrypoint': 'gsm8k_app.py',
    'host': '127.0.0.1',
    'port': 18575,
    'pid': 1582343,
    'config_path': 'gsm8k_aviary_resources_server',
    'url': 'http://127.0.0.1:18575',
}
[2] gsm8k_aviary_agent (responses_api_agents/aviary_agent)
{
    'process_name': 'gsm8k_aviary_agent',
    'server_type': 'responses_api_agents',
    'name': 'aviary_agent',
    'dir_path': (
        '/home/ubuntu/Gym/responses_api_agents/aviary_agent'
    ),
    'entrypoint': 'app.py',
    'host': '127.0.0.1',
    'port': 63115,
    'pid': 1582344,
    'config_path': 'gsm8k_aviary_agent',
    'url': 'http://127.0.0.1:63115',
}
[3] policy_model (responses_api_models/vllm_model)
{
    'process_name': 'policy_model',
    'server_type': 'responses_api_models',
    'name': 'vllm_model',
    'dir_path': (
        '/home/ubuntu/Gym/responses_api_models/vllm_model'
    ),
    'entrypoint': 'app.py',
    'host': '127.0.0.1',
    'port': 55951,
    'pid': 1582347,
    'config_path': 'policy_model',
    'url': 'http://127.0.0.1:55951',
}
####################################################################################################

接下来,在 GSM8K 环境中运行智能体。以下命令将针对输入文件中的 5 个示例问题执行简单代理,并将代理的运行轨迹记录到输出文件中。

ng_collect_rollouts \
    +agent_name=gsm8k_aviary_agent \
    +input_jsonl_fpath=resources_servers/aviary/data/gsm8k_example.jsonl \
    +output_jsonl_fpath=results/gsm8k_aviary_rollouts.jsonl


您应该会收到一个显示轨迹平均奖励的输出:

Collecting rollouts: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:18<00:00,  3.71s/it]
{
    "reward": 1.0,
}

为了查看轨迹,NeMo Gym 提供了一个简洁的用户界面:

ng_viewer +jsonl_fpath=results/gsm8k_aviary_rollouts.jsonl

第 4 步:构建新环境

要在 NeMo Gym 中创建新环境,您可以先在 Aviary 中构建环境,再通过 Aviary 集成便捷地创建新的资源服务器;或者,您也可以直接在 NeMo Gym 中从头开始构建自定义环境。在本示例中,我们将 Aviary 的 HotPotQA 环境添加到 NeMo Gym。首先,创建 `resources_servers/aviary/hotpotqa_app.py` 文件,以扩展基础的飞行员资源服务器:
First, create `resources_servers/aviary/hotpotqa_app.py`,which extends the base aviary resources server:

from pydantic import Field
from aviary.envs.hotpotqa import HotPotQADataset, HotPotQAEnv
from resources_servers.aviary.app import AviaryResourcesServer


class HotPotQAResourcesServer(AviaryResourcesServer[HotPotQAEnv, HotPotQADataset]):
   dataset: HotPotQADataset = Field(default_factory=lambda: HotPotQADataset(split="train"))


if __name__ == "__main__":
   HotPotQAResourcesServer.run_webserver()

接下来,在 `resources_servers/aviary/configs/hotpotqa_aviary.yaml` 中创建配置文件:

hotpotqa_aviary_resources_server:
  resources_servers:
    aviary:
      entrypoint: hotpotqa_app.py
hotpotqa_aviary_agent:
  responses_api_agents:
    aviary_agent:
      entrypoint: app.py
      resources_server:
        type: resources_servers
        name: hotpotqa_aviary_resources_server
      model_server:
        type: responses_api_models
        name: policy_model
      datasets:
      - name: train
        type: train
        jsonl_fpath: resources_servers/aviary/data/hotpotqa_train.jsonl
        gitlab_identifier:
          dataset_name: hotpotqa_train
          version: 0.0.1
          artifact_fpath: hotpotqa_train.jsonl
        license: Apache 2.0
      - name: validation
        type: validation
        jsonl_fpath: resources_servers/aviary/data/hotpotqa_validation.jsonl
        gitlab_identifier:
          dataset_name: hotpotqa_validation
          version: 0.0.1
          artifact_fpath: hotpotqa_validation.jsonl
        license: Apache 2.0
      - name: hotpotqa_example
        type: example
        jsonl_fpath: resources_servers/aviary/data/hotpotqa_example.jsonl
        gitlab_identifier:
          dataset_name: hotpotqa_example
          version: 0.0.1
          artifact_fpath: hotpotqa_example.jsonl
        license: Apache 2.0

然后,在 `resources_servers/aviary/data/hotpotqa_example.jsonl` 中创建一个示例数据集,该数据集提供任务索引,用于从底层 Aviary 环境数据集中检索样本:

{"task_idx":0,"responses_create_params":{"input":[]}}
{"task_idx":1,"responses_create_params":{"input":[]}}
{"task_idx":2,"responses_create_params":{"input":[]}}
{"task_idx":3,"responses_create_params":{"input":[]}}
{"task_idx":4,"responses_create_params":{"input":[]}}

随后,更新 `requirements.txt` 以包含 Aviary 中的 `hotpotqa` 包:

-e nemo-gym[dev] @ ../../
fhaviary[gsm8k,hotpotqa,notebook,llm]>=0.24.1
tqdm
datasets
huggingface-hub

通过这四项修改,我们现在可以在 NeMo Gym 中运行 Aviary HotPotQA 环境。

访问 NeMo Gym 资源库,了解更多即用型训练环境。产品文档还全面介绍了关键概念、资源服务器的创建方法以及 RL 训练的执行步骤。查看我们最新发布的 NeMo RL 版本,支持 on-policy 蒸馏、asyncRL、高级 RL 算法和端到端 FP8 RL 训练。

构建科学智能体的最佳实践

构建科学智能体具有挑战性,但通过采取以下实践,团队可以在打造功能更强大的系统方面稳步取得进展。

  • 从简单着手。从基本的智能体开始,而非直接构建多智能体系统或集成大量工具。在引入复杂的奖励机制之前,优先采用基于结果的奖励机制,以避免奖励被利用或操纵。
  • 进行奖励分析。当模型能够为任务生成多种解决方案(其中部分方案正确)时,采用 GRPO 风格的算法进行训练效果显著。通过多次采样并计算每个任务的奖励均值与标准差,有助于构建更高效的训练环境。
  • 监控训练指标。各类反映训练稳定性、模型行为及学习进展的指标将自动记录至 Weights & Biases。例如,通过分析这些指标,可及时发现采样偏差、模型崩溃或轨迹截断等问题。
  • 延长训练时间。使用基于 RLVR 的方法训练时,初期学习进展通常较缓慢,而在后期阶段学习曲线会明显上升。这种情况常出现在模型初期难以找到正确解法,但随着训练深入逐渐掌握有效策略的过程中。

这些步骤为利用 NeMo Gym、NeMo RL 和 Aviary 大规模构建与训练科学智能体提供了切实可行的路径。立即着手开发您的新一代科学智能体。如需了解更多信息,请访问以下链接:

此外,还可查看我们全新的 NVIDIA Nemotron 3 模型系列,该系列现已支持 Nano。

参与此项工作的人员包括 Brian Yu、Chris Wing、Elliot Eshelman 和 Sylendran Arunagiri。

 

 

标签