使用 Agent-lightning 和 VERL 训练 SQL Agent¶

本教程基于 Agent-lightning SQL Agent 示例，并解释了系统组件如何集成：一个作为 LitAgent 包装的 基于 LangGraph 的 SQL agent，VERL 强化学习 (RL) 算法，以及 Trainer，它协调训练和调试。

在 examples/spider/train_sql_agent.py 中的命令行界面提供了一个完整的可运行示例。但是，本文档侧重于理解底层架构，以便您可以有效地将工作流程调整到您自己的 agent。

SQL Agent 架构¶

Agent-lightning 可以与各种编排框架无缝集成，包括 Agent Framework、AutoGen、CrewAI、LangGraph 和 OpenAI Agents SDK。它还可以与自定义 Python 逻辑互操作。

在本示例中，LangGraph 定义了一个循环工作流程，该工作流程模拟了分析师迭代 SQL 开发过程。下图（直接从 sql_agent.py 渲染）说明了 agent 如何起草、执行、评论和完善查询，直到获得令人满意的结果。

---
config:
  flowchart:
    curve: linear
---
graph LR;
        __start__([<p>__start__</p>]):::first
        write_query(write_query)
        execute_query(execute_query)
        check_query(check_query)
        rewrite_query(rewrite_query)
        __end__([<p>__end__</p>]):::last
        __start__ --> write_query;
        check_query -.-> __end__;
        check_query -.-> rewrite_query;
        execute_query --> check_query;
        rewrite_query --> execute_query;
        write_query --> execute_query;
        classDef default fill:#f2f2f2,line-height:1.2
        classDef first fill-opacity:0
        classDef last fill:#cccccc

注意

工作流程通过以下阶段进行

write_query – 根据用户的提问和数据库模式生成初始 SQL 查询。
execute_query – 在目标数据库上执行生成的查询。
check_query – 使用专门的提示（CHECK_QUERY_PROMPT）评估查询及其结果（或错误），以检测问题。
rewrite_query – 如果发现问题，agent 将使用来自上一步的反馈重写查询并重新进入循环。
END – 当查询被验证或达到最大迭代次数（max_turns）时，循环终止。每个 turn 包含一个完整的 write_query、execute_query、check_query 和（如果适用）rewrite_query 阶段的循环。

在本教程中，强化学习 (RL) 用于优化 write_query 和 rewrite_query 阶段。虽然 check_query 步骤共享相同的底层 LLM 权重，但其跟踪数据不用于学习。

为了保持设计模块化和可维护性，建议将基于 LangGraph 的 SQL Agent 定义在一个单独的文件中，并通过一个构建器函数公开，例如

def build_langgraph_sql_agent(
    database_path: str,
    openai_base_url: str,
    model: str,
    sampling_parameters: Dict[str, Any],
    max_turns: int,
    truncate_length: int
):
    builder = StateGraph(State)
    builder.add_node(write_query)
    ...

    builder.add_edge(START, "write_query")
    ...

    return builder.compile().graph()

这种方法将您的 LangGraph 逻辑与 Agent-lightning 版本更改隔离，从而提高可读性和可调试性。

桥接 LangGraph 和 Agent-lightning¶

提示

在阅读本节时，请打开 sql_agent.py。这将帮助您理解此处显示的片段的工作原理。

在 sql_agent.py 中定义的 LitSQLAgent 类充当桥梁。它继承自 agl.LitAgent，允许 runner 为每个 rollout 提供共享资源（例如 LLM）。

以下是一个简化的逻辑说明（请注意：这是一种概念性伪代码；实际实现包括数据集特定的细节）

class LitSQLAgent(agl.LitAgent[Dict[str, Any]]):

    def __init__(self, max_turns: int, truncate_length: int):
        # Every turn here refers to a full cycle of write/exe/check/rewrite
        self.max_turns = max_turns
        self.truncate_length = truncate_length

    def rollout(
        self,
        task: Dict[str, Any],
        resources: agl.NamedResources,
        rollout: agl.Rollout
    ) -> float | None:
        llm: agl.LLM = resources["main_llm"]
        agent = build_langgraph_sql_agent(
            database_path="sqlite:///" + task["db_id"],
            max_turns=self.max_turns,
            truncate_length=self.truncate_length,
            openai_base_url=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
            model=llm.model,
            sampling_parameters=llm.sampling_parameters,
        )
        result = agent.invoke({"question": question}, {
            "callbacks": [self.tracer.get_langchain_handler()],
            "recursion_limit": 100,
        })
        reward = evaluate_query(result["query"], ground_truth, db_path, raise_on_error=False)
        return reward

LitSQLAgent 作为 LangGraph agent 周围的一个轻量级包装器，提供正确的接口给 rollout 方法。它构建 LangGraph agent，调用它，并将评估结果作为奖励信号返回。

"main_llm" 资源键是 agent 和 VERL 之间的约定。它用于在 rollout 期间从 VERL 算法注入一个 OpenAI 兼容的端点。支持两种使用此 agentlightning.LLM 资源的方法

直接访问 – 使用 llm.endpoint 进行简单的集成（与 v0.1 示例相同）。
感知上下文的访问 – 使用 get_base_url 与 rollout.rollout_id 和 rollout.attempt.attempt_id。这种方法可以实现每调用者跟踪归因，从而在 runner 端跟踪器不可用时提高每个 rollout 或尝试的跟踪收集。有关详细信息，请参阅使用 Traces。

奖励信号和评估¶

evaluate_query 函数为 RL 训练提供奖励机制。在 agent 训练中，获得一致且有意义的奖励信号通常具有挑战性。幸运的是，当使用 Spider 数据集 时，这得到了简化。该数据集包含约 8k 个样本，包含自然语言问题、数据库模式和真实的 SQL 查询。

使用 Spider 评估器，agent 生成的查询在目标数据库上执行，并与真实的查询进行比较。如果两个查询产生相同的执行结果，则认为它们是等效的。

注意

必须绝不在训练期间向 agent 暴露真实的查询，以防止数据泄漏。

在此设置中，奖励直接从 rollout 方法返回，使 runner 能够将其转发回 RL 算法。

警告

避免同时使用 emit_reward 和返回奖励值。这样做会导致算法接收到重复的奖励信号，从而导致训练行为不一致。

配置 VERL 用于强化学习¶

请查看 examples/spider/train_sql_agent.py 以获取完整的强化学习配置，这是一个纯 Python 字典。它镜像（实际上是）shell 参数，用于启动 VERL 框架中的训练，但更易于以编程方式调整

verl_config: Dict[str, Any] = {
    "algorithm": {"adv_estimator": "grpo", "use_kl_in_reward": False},
    "data": {
        # train_files and val_files are no longer needed here
        # because data are read in agl.Trainer
        ...,
        # Controls how many tasks are pooled per step
        # (multiplied by actor_rollout_ref.rollout.n)
        "train_batch_size": 32,
        # Prompt and responses larger than these lengths are truncated
        "max_prompt_length": 4096,
        "max_response_length": 2048,
    },
    "actor_rollout_ref": {
        "rollout": {
            # Only vLLM is supported currently
            "name": "vllm",
            # Equals to group size of GRPO
            "n": 4,
            # Used to enable tool call parser in vLLM
            "multi_turn": {"format": "hermes"},
            ...
        },
        "actor": {"ppo_mini_batch_size": 32, "optim": {"lr": 1e-6}, ...},
        "model": {
            # Config your preferred LLM here
            "path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
            ...
        },
    },
    "trainer": {
        "n_gpus_per_node": 1,
        # Validation once before training starts
        "val_before_train": True,
        # Validation every N training steps
        "test_freq": 32,
        # Save checkpoints every N training steps
        "save_freq": 64,
        # Go through the train dataset this many times
        "total_epochs": 2
    },
}

这等效于以下 CLI 调用

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    algorithm.use_kl_in_reward=False \
    data.train_batch_size=32 \
    data.max_prompt_length=4096 \
    data.max_response_length=2048 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.n=4 \
    actor_rollout_ref.rollout.multi_turn.format=hermes \
    actor_rollout_ref.actor.ppo_mini_batch_size=32 \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-Coder-1.5B-Instruct \
    trainer.n_gpus_per_node=1 \
    trainer.val_before_train=True \
    trainer.test_freq=32 \
    trainer.save_freq=64 \
    trainer.total_epochs=2

警告

我们曾经在 v0.1 中提供一个名为 python -m agentlightning.verl 的 CLI 来启动训练。这不再是推荐的方法。相反，请使用 agl.Trainer 来一起运行 VERL 和 agent runner，或者按照调试教程，如果您想要类似于 v0.1 的隔离体验。

使用 Trainer 编排训练¶

Trainer 是高级编排器，它集成了 agent、算法、数据集和分布式 runner。使用 Trainer 的主要好处是

它允许您使用一行代码启动所有内容：trainer.fit(...)。
它暴露了诸如 n_runners 之类的配置选项来控制并行度，以及 adapter 来定义算法如何解释代理产生的跟踪数据。

以下是一个示例用法

import agentlightning as agl

agent = LitSQLAgent()
algorithm = agl.VERL(verl_config)
trainer = agl.Trainer(
    n_runners=10,
    algorithm=algorithm,
    adapter={"agent_match": active_agent},
)
train_data = pd.read_parquet("data/train_spider.parquet").to_dict("records")
val_data = pd.read_parquet("data/test_dev_500.parquet").to_dict("records")
trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)

首先，agl.VERL(verl_config) 启动 VERL 算法及其 OpenAI 兼容的代理。train_data 和 val_data 被传递到 VERL，它将任务排队到一个由 LightningStore 管理的集中式任务队列，所有 runner 都可以访问。

当调用 Trainer.fit 时，它会启动 10 个并发 runner（如 n_runners=10 所指定）。每个 runner 从集中式任务队列中拉取任务，执行代理的 rollout 方法，收集轨迹，并将奖励返回给 VERL 进行训练。

如前所述，Adapter 在算法侧使用，接收代理和 runner 产生的轨迹。agent_match 参数确保 VERL 仅摄取来自您想要优化的特定代理的 span。在上面的示例中，至少有三个代理——write_query、rewrite_query 和 check_query。通过将 agent_match 设置为像 "write" 这样的正则表达式，可以同时优化 write_query 和 rewrite_query 代理。如果需要，您也可以将其设置为 "write|check" 或 None 以包含所有代理。

使用 `Trainer.dev` 对 Pipeline 进行 Dry-Run¶

在投入数小时的 GPU 时间之前，您可以使用 Trainer.dev() 对代理进行 dry-run。此方法切换到轻量级的 Baseline 算法，排队最多十个任务，并打印代理发出的每个 span。因为它使用与完整训练相同的 runner 堆栈，因此非常适合验证数据库连接和 LangGraph 控制流。

首先，由于 VERL 在此模式下不处于活动状态，代理需要一个有效的 OpenAI 兼容的端点。您可以使用 OpenAI 的官方 API 或您自己的本地 LLM 端点。将其包装如下

trainer = agl.Trainer(
    n_workers=1,
    initial_resources={
        "main_llm": agl.LLM(
            endpoint=os.environ["OPENAI_API_BASE"],
            model="gpt-4.1-nano",
            sampling_parameters={"temperature": 0.7},
        )
    },
)

然后，使用少量任务调用 trainer.dev(...)

dev_data = pd.read_parquet("data/test_dev_500.parquet").to_dict("records")[:10]
trainer.dev(agent, dev_dataset=dev_data)

在 Python 会话中运行此操作，或调整您的脚本以包含 --dev 标志。一旦 span 看起来正常并且奖励不为零，就切换回 trainer.fit(...) 进行完整的 RL 训练。有关调试代理的更多技巧，请参阅调试教程。

运行示例代码¶

以下教程解释了如何在 examples/spider 中运行完整的示例。

数据集¶

trainer 期望在 examples/spider/data 目录中包含三个 Parquet 文件：train_spider.parquet、test_dev_500.parquet 和 test_dev.parquet。

下载存储库提供的策划数据集包

cd examples/spider
pip install gdown  # included in the 'experiment' optional dependency
gdown --fuzzy https://drive.google.com/file/d/1oi9J1jZP9TyM35L85CL3qeGWl2jqlnL6/view
unzip -q spider-data.zip -d data
rm spider-data.zip

如果您更喜欢自己生成文件，请下载 Spider 1.0 并运行

python spider_eval/convert_dataset.py

如果您将数据集存储在默认 data 目录之外，请设置 VERL_SPIDER_DATA_DIR。

依赖项¶

创建一个干净的虚拟环境，激活它，并使用本教程所需的 VERL extras 安装 Agent-lightning。根据需要安装 LangChain 相关依赖项。

对于完整的训练配置，计划使用至少 40 GB 内存的 GPU。

启动训练¶

从 examples/spider，根据您的模型偏好运行以下助手脚本

python train_sql_agent.py qwen   # Default Qwen-2.5-Coder-1.5B run
python train_sql_agent.py llama  # LLaMA-3.2-1B with llama3_json tool parser

该脚本实例化 LitSQLAgent 并启动 trainer.fit。如果您只想训练图中的一个代理，请提供 --active-agent my_agent_variant。

对于 LLaMA 配置，在运行之前导出 HF_TOKEN，以便 VERL 可以下载模型权重。

故障排除

如果您在 WANDB_API_KEY 未设置、HF_TOKEN 未设置或数据未找到时遇到 Ray worker 错误，请尝试使用助手脚本重新启动 Ray 集群：scripts/restart_ray.sh，这本质上会停止任何 Ray 集群并启动一个新的集群

env RAY_DEBUG=legacy HYDRA_FULL_ERROR=1 VLLM_USE_V1=1 ray start --head --dashboard-host=0.0.0.0

使用 NPU 启动训练

该示例还支持使用 华为昇腾 NPU 运行。此功能由华为团队贡献。要使用它，请使用脚本中的 config_train_npu 函数。

支持的硬件：Atlas 200T A2 Box16、Atlas 900 A2 PODc、Atlas 800T A3。运行 Qwen2.5-Coder-1.5B-Instruct 模型至少需要 单个 40GB NPU。

环境设置：Python 3.11.13、CANN 8.2.RC1、torch 2.7.1+cpu、torch_npu 2.7.1.dev20250724。有关基本的环境准备，请参阅此文档。

在安装依赖项之前，配置以下 pip 镜像

pip config set global.index-url http://repo.huaweicloud.com/repository/pypi/simple
pip config set global.extra-index-url "https://download.pytorch.org/whl/cpu/ https://mirrors.huaweicloud.com/ascend/repos/pypi"

然后安装 vLLM、vLLM-Ascend 和 VERL

pip install vllm==0.10.0 --trusted-host repo.huaweicloud.com
pip install vllm-Ascend==0.10.0rc1 --trusted-host repo.huaweicloud.com
pip install verl==0.5.0

为了确保 VERL 框架在 NPU 上正确运行，请将以下行添加到 verl/utils/vllm_utils.py

from vllm_ascend.patch import platform
from vllm_ascend.patch import worker

有关更多详细信息，请参阅以下参考：https://github.com/vllm-project/vllm-ascend/issues/1776。

安装完上述依赖项后，从 examples/spider 运行以下脚本命令

python train_sql_agent.py npu

在没有 VERL 的情况下调试代理¶

sql_agent.py 还提供了一个 debug_sql_agent() 助手，用于在将 LangGraph 工作流直接用于本地或托管的 OpenAI 兼容端点之前使用 VERL。

设置以下环境变量，然后执行该文件

export OPENAI_API_BASE=<your_api_base>
export OPENAI_API_KEY=<your_api_key>
cd examples/spider
python sql_agent.py

这使您可以在引入强化学习之前验证工作流和提示的行为是否符合预期。

评估¶

以下结果是通过在单个 80 GB GPU 上运行 python train_sql_agent.py qwen 获得的。训练大约需要 12 小时。下面的训练曲线通过每 16 步进行聚合以获得更好的可视化效果。

使用旧版本 — Agent-lightning v0.1.1、verl==0.5.0 和 vllm==0.10.0 收集了其他评估结果。您可以在此文章中找到它们：使用强化学习训练 AI 代理来编写和自我纠正 SQL