恒美微站
首页
关于我们
建站服务
主题模板
案例展示
资讯中心
联系我们
slime × Tau-Bench 实战:Agentic 多轮工具调用环境下的 RL 训练指南
首页
资讯中心
/
slime × Tau-Bench 实战:Agentic 多轮工具调用环境下的 RL 训练指南
slime × Tau-Bench 实战:Agentic 多轮工具调用环境下的 RL 训练指南
发布时间:2026/9/16 20:48:24
slime × Tau-Bench 实战Agentic 多轮工具调用环境下的 RL 训练指南【免费下载链接】slimeslime is an LLM post-training framework for RL Scaling.项目地址: https://gitcode.com/GitHub_Trending/slime12/slime本文围绕 slime 开源仓库中examples/tau-bench示例讲解如何在 tau-benchretail / airline 客服对话模拟环境中完成 Agentic 多轮工具调用数据的构造、模型 checkpoint 准备并通过 slime 自定义生成函数--custom-generate-function-path打通环境交互 → 轨迹采样 → GRPO 训练的完整闭环。读完本文你将掌握 tau-bench 与 slime 集成的全部配置细节、核心源码调用链以及如何改造为自己的多轮 Agent RL 训练方案。一、示例定位为什么在 slime 中引入 Tau-Benchslime 是一个面向 RL Scaling 的 LLM 后训练框架。其默认 rollout 路径是给定 prompt一次性生成完整回复而真实世界的 Agent 任务需要多轮决策模型每轮输出一个动作调用工具或直接回复环境返回观察结果模型再基于新状态继续决策直到任务结束。tau-bench 提供了标准化的retail零售客服与airline航空客服两个仿真环境内置了用户模拟器、工具集与奖励判定非常适合用来验证 Agent 的多轮工具调用能力。examples/tau-bench目录即为此设计其完整文件清单如下文件作用tau1_mock.py生成 slime 训练所需的 JSONL 任务索引数据generate_with_tau.pyslime 自定义生成函数入口负责一次完整的 Agent-环境交互trainable_agents.py可训练的 ToolCallingAgent异步调用 sglang、多轮循环、token 统计token_delta.py多轮对话的 token 增量与 loss mask 计算openai_tool_adapter.py把 sglang 工具调用解析结果转换为 OpenAI 兼容格式sglang_tool_parser.py复用 sglang 的FunctionCallParser做本地工具解析run_qwen3_4B.sh一键启动 2 卡 slime 训练脚本对应测试见 test_tau_bench_token_delta.py用于验证多轮 token 增量计算的正确性。二、环境准备与依赖安装2.1 镜像与基础环境使用zhuzilin/slime:latest镜像并按以下步骤初始化 slime 本体及 tau-bench 依赖注意 tau-bench 需要切换到feature/litellm-retry分支以获得用户模拟器调用时的重试能力cd /root/ git clone https://github.com/THUDM/slime.git cd slime pip install -e . --no-deps # for tau bench cd /root/ git clone https://github.com/JD-ETH/tau-bench.git cd tau-bench git checkout feature/litellm-retry pip install -e . --no-depsslime 采用可编辑安装pip install -e . --no-deps依赖按 requirements.txt 另行安装tau-bench 同样以可编辑方式安装便于后续直接调用其tau_bench.envs、tau_bench.types等模块。2.2 生成训练用 mock 数据tau-bench 任务数据通过环境实例暴露tau1_mock.py会遍历{retail: [train, test, dev], airline: [test]}的所有组合把每个任务索引与元数据写为 JSONL一行一个 JSON 对象供 slime 按索引回放cd /root/slime/examples/tau-bench python tau1_mock.py --local_dir /root/tau-bench/生成的数据文件包括retail_train_tasks.jsonl、retail_dev_tasks.jsonl、retail_test_tasks.jsonl与airline_test_tasks.jsonl。每行形如{index: 0, metadata: {task_id: ..., display_name: ...}}其实现tau1_mock.py核心逻辑如下构造RunConfig(model_providermock, user_model_providermock, user_strategyhuman, modelmock)后调用get_env(env_name..., user_strategyhuman, ...)获取环境实例并读取其tasks属性ALL_DATA_MAPPINGS {retail: [train, test, dev], airline: [test]} # ... for env, split in ALL_DATA_MAPPINGS.items(): for s in split: config.env env config.task_split s env_instance get_env( env_nameconfig.env, user_strategyconfig.user_strategy, user_modelconfig.user_model, task_splitconfig.task_split, ) output_path os.path.join(local_dir, f{env}_{s}_tasks.jsonl) with open(output_path, w) as f: for i, task in enumerate(env_instance.tasks): row {index: i, metadata: task.model_dump()} f.write(json.dumps(row) \n)注意user_strategyhuman与model_providermock仅用于离线枚举任务不会真正调用用户模拟器。2.3 准备 Qwen3-4B-Instruct-2507 的两种 checkpoint训练使用 Qwen3-4B-Instruct-2507后续模型参数可在 scripts/models/qwen3-4B-Instruct-2507.sh 中按需调整。它通过MODEL_ARGS_ROTARY_BASE5000000继承 scripts/models/qwen3-4B.sh 的架构参数并覆盖了 RoPE base 为 5000000MODEL_ARGS_ROTARY_BASE5000000 source $(dirname -- ${BASH_SOURCE[0]})/qwen3-4B.sh分别准备 HuggingFace 原始权重与 Megatron-Coremcore分布式权重# hf checkpoint hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen3-4B-Instruct-2507 # mcore checkpoint cd /root/slime source scripts/models/qwen3-4B-Instruct-2507.sh PYTHONPATH/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ ${MODEL_ARGS[]} \ --hf-checkpoint /root/Qwen3-4B-Instruct-2507 \ --save /root/Qwen3-4B-Instruct-2507_torch_dist其中tools/convert_hf_to_torch_dist.py见 tools/convert_hf_to_torch_dist.py把 HF 权重转换为 mcore 的 torch 分布式格式作为训练/回放阶段--ref-load与--load的初始来源。三、配置 TAU_CONFIGS 与用户模拟器 APIgenerate_with_tau.py顶部的TAU_CONFIGS是整套集成的关键配置运行时会被转换为RunConfig传给 tau-benchTAU_CONFIGS { env: retail, # Select between [retail, airline] agent_strategy: tool-calling, # Select between [tool-calling, act, react, few-shot], only tool-calling implemented for now user_model: gemini-2.0-flash-lite, # Cheap Model for user simulator user_model_provider: gemini, task_split: train, # Select between [train, test, dev] for retail, [test] for airline user_strategy: llm, # Select between [llm, react, verify, reflection] model_provider: auto_router, # Unused, required model: qwen3-4b, # Unused, required } # Replace with your actual API key for user sim GEMINI_API_KEY YOUR KEY各字段含义与取值范围配置项可选值说明envretail/airline选择仿真环境agent_strategytool-calling其余act/react/few-shot暂未实现策略智能体行为源码中agent_factory对非tool-calling直接抛出NotImplementedErroruser_model如gemini-2.5-flash-lite用户模拟器使用的廉价模型user_model_providergemini等用户模拟器模型供应商task_splittrain/test/devairline 仅test任务划分user_strategyllm/react/verify/reflection用户模拟器行为策略model_provider/model任意占位值代码中未使用但为满足RunConfig必填字段注意GEMINI_API_KEY会被写入环境变量os.environ[GEMINI_API_KEY]供 tau-bench 的用户模拟器调用外部 LLM 时使用。请务必替换为你自己的有效 Key否则用户模拟环节会失败。四、一键启动训练run_qwen3_4B.sh 参数全解配置完成后运行cd /root/slime bash examples/tau-bench/run_qwen3_4B.sh该脚本run_qwen3_4B.sh会先清理残留进程pkill sglang、ray stop --force等然后以 Ray 作业方式提交train.py。其参数按功能拆分为多组下面逐一解读。4.1 清理与基础环境pkill -9 sglang sleep 3 ray stop --force pkill -9 ray pkill -9 python sleep 3 ... export PYTHONUNBUFFERED1PYTHONUNBUFFERED1防止 Ray 缓冲 stdout/stderr便于实时观察日志。脚本还会探测 NVLink 数量并打印HAS_NVLINK用于后续性能相关判断。4.2 checkpoint 参数CKPT_ARGSCKPT_ARGS( --hf-checkpoint /root/Qwen3-4B-Instruct-2507/ --ref-load /root/Qwen3-4B-Instruct-2507_torch_dist/ --load /root/Qwen3-4B-Instruct-2507_slime/ --save /root/Qwen3-4B-Instruct-2507_slime/ --save-interval 20 )--hf-checkpointHF 原始权重用于加载 tokenizer 与处理器--ref-loadmcore 参考模型权重用于 KL 等计算--load/--save训练模型的加载/保存目录--save-interval 20每 20 步保存一次 checkpoint。4.3 rollout 数据与采样ROLLOUT_ARGSROLLOUT_ARGS( --prompt-data /root/tau-bench/retail_train_tasks.jsonl --input-key index --rollout-shuffle --num-rollout 500 --rollout-batch-size 32 --n-samples-per-prompt 8 --rollout-max-response-len 1024 --rollout-temperature 1 --global-batch-size 256 --dynamic-sampling-filter-path slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std --balance-data )--prompt-data指向tau1_mock.py生成的retail_train_tasks.jsonl--input-key index说明generate_with_tau.py中int(sample.prompt)取出的正是每行的index字段用于精确回放任务--num-rollout 500、--rollout-batch-size 32、--n-samples-per-prompt 8共 500 次 rollout每批 32 条 prompt每条 prompt 采样 8 个轨迹--global-batch-size 256训练全局 batch 大小8 × 32--rollout-max-response-len 1024多轮轨迹总长度上限--dynamic-sampling-filter-path指定按同组奖励标准差非零过滤样本即check_reward_nonzero_std。其实现dynamic_sampling_filters.py为def check_reward_nonzero_std(args, samples: list[Sample], **kwargs): rewards [sample.get_reward_value(args) for sample in samples] keep torch.tensor(rewards, dtypetorch.float64).std() 1e-6 # ... 若 keep 为 False 则过滤整组即同一 prompt 的 8 条轨迹若奖励完全相同std ≤ 1e-6则整组被剔除保证 GRPO 的优势估计有区分度--balance-data对数据做均衡采样。4.4 评估参数EVAL_ARGSEVAL_ARGS( --eval-interval 5 --eval-prompt-data retail-dev /root/tau-bench/retail_dev_tasks.jsonl --n-samples-per-eval-prompt 1 --eval-max-response-len 1024 --eval-top-k 1 )每 5 个训练步在retail_dev_tasks.jsonl上做一次贪心--eval-top-k 1评估观察验证集上的工具调用成功率与奖励变化。4.5 并行与性能PERF_ARGSPERF_ARGS( --tensor-model-parallel-size 2 --sequence-parallel --pipeline-model-parallel-size 1 --context-parallel-size 1 --expert-model-parallel-size 1 --expert-tensor-parallel-size 1 --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 --use-dynamic-batch-size --max-tokens-per-gpu 9216 )本示例在 2 卡上运行使用 TP2 序列并行--recompute-granularity full --recompute-method uniform --recompute-num-layers 1用激活重计算节省显存--use-dynamic-batch-size --max-tokens-per-gpu 9216让 sglang 依据每卡 token 上限动态调整 batch。4.6 GRPO 与优化器GRPO_ARGS / OPTIMIZER_ARGSGRPO_ARGS( --advantage-estimator grpo --use-kl-loss --kl-loss-coef 0.00 --kl-loss-type low_var_kl --entropy-coef 0.00 --eps-clip 0.2 --eps-clip-high 0.28 ) OPTIMIZER_ARGS( --optimizer adam --lr 1e-6 --lr-decay-style constant --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 )采用 GRPO 优势估计器eps-clip 0.2 / eps-clip-high 0.28设置裁剪上下界KL 系数置 0纯奖励优化学习率 1e-6 恒定衰减Adam 优化器。4.7 sglang 服务、杂项与自定义生成函数SGLANG_ARGS( --rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.7 # If gemini API reports concurrency limit error, set this parameter to reduce the concurrency # --sglang-server-concurrency 32 ) MISC_ARGS( --attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 --attention-backend flash ) CUSTOM_ARGS( --custom-generate-function-path generate_with_tau.generate )sglang 每引擎 1 卡、静态显存 0.7若 Gemini API 报并发限制错误可调低--sglang-server-concurrency关闭 dropout、fp32 累积梯度与 softmax、flash attention 后端注意使用 MLA 架构的模型时需注释--attention-backend flash核心--custom-generate-function-path generate_with_tau.generate把默认的单轮生成替换为 tau-bench 多轮交互生成。4.8 Ray 集群与作业提交export MASTER_ADDR${MASTER_ADDR:-127.0.0.1} NUM_GPUS2 ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} \ --disable-usage-stats --dashboard-host0.0.0.0 --dashboard-port8265 --temp-dir /root/shared/ray_temp RUNTIME_ENV_JSON{ env_vars: { PYTHONPATH: /root/Megatron-LM/:${SCRIPT_DIR}, CUDA_DEVICE_MAX_CONNECTIONS: 1 } } ray job submit --addresshttp://127.0.0.1:8265 \ --runtime-env-json${RUNTIME_ENV_JSON} \ -- python3 train.py \ --actor-num-nodes 1 \ --actor-num-gpus-per-node ${NUM_GPUS} \ --rollout-num-gpus ${NUM_GPUS} \ --colocate \ ${MODEL_ARGS[]} ${CKPT_ARGS[]} ${ROLLOUT_ARGS[]} \ ${OPTIMIZER_ARGS[]} ${GRPO_ARGS[]} ${DISTRIBUTED_ARGS[]} \ ${WANDB_ARGS[]} ${PERF_ARGS[]} ${EVAL_ARGS[]} \ ${SGLANG_ARGS[]} ${MISC_ARGS[]} ${CUSTOM_ARGS[]}--colocate表示 actor 与 rollout 引擎共用同一批 GPUPYTHONPATH注入/root/Megatron-LM/与脚本目录保证generate_with_tau等自定义模块可被 Ray worker 导入。五、源码拆解多轮交互如何进入 slime 训练管线5.1 自定义生成函数的挂载点在 slime 的 rollout 引擎 sglang_rollout.py 中generate_and_rm会检查--custom-generate-function-path并动态加载custom_func_path getattr(sample, generate_function_path, None) or args.custom_generate_function_path if custom_func_path is not None: custom_generate_func load_function(custom_func_path) # if signature has evaluation, pass evaluation if evaluation in inspect.signature(custom_generate_func).parameters: sample await custom_generate_func(args, sample, sampling_params, evaluationevaluation) else: sample await custom_generate_func(args, sample, sampling_params) else: sample await generate(args, sample, sampling_params)因此自定义函数只需满足async def generate(args, sample, sampling_params) - Sample签名即可被无缝调用若函数内部已填充sample.reward后续async_rm会自动跳过奖励模型。同时GenerateStatesglang_rollout.py是单例缓存 tokenizer 与sampling_params的默认值temperature、top_p、top_k、max_new_tokens、stop 等供自定义函数内复用。5.2 generate_with_tau.generate一次完整 Agent-环境交互generate_with_tau.py 是集成的总入口其流程如下async def generate(args, sample, sampling_params) - Sample: # 不支持 partial rollout assert not args.partial_rollout, Partial rollout is not supported for tau-bench interactions. # 从 sample.prompt 中提取任务索引对应 --input-key index task_index int(sample.prompt) # 初始化 tau-bench 环境 env get_env( env_nametau_config.env, user_strategytau_config.user_strategy, user_modeltau_config.user_model, user_providertau_config.user_model_provider, task_splittau_config.task_split, task_indextask_index, ) # 创建可训练智能体sglang 异步交互 agent agent_factory( tools_infoenv.tools_info, wikienv.wiki, configtau_config, rollout_argsargs, sampling_paramssampling_params, ) # 执行多轮交互 interaction_result await agent.asolve(env, agent.rollout_args, agent.sampling_params, task_index) # 转换为 slime Sample result_sample res_to_sample(interaction_result, task_index) return result_sample其中res_to_sample完成状态映射与response_length兜底计算status_mapping { Status.COMPLETED: completed, Status.TRUNCATED: truncated, Status.ABORTED: aborted, } status status_mapping.get(res.status) sample Sample( indextask_index, promptres.prompt, tokensres.tokens, responseres.response, rewardres.reward, loss_maskres.loss_mask, statusstatus, metadatares.info, ) if hasattr(res, response_length): sample.response_length res.response_length elif res.loss_mask: sample.response_length len(res.loss_mask) # loss_mask 只覆盖回复部分 elif res.tokens: sample.response_length len(res.tokens) else: sample.response_length 0Sample定义于 slime/utils/types.py其response_length与loss_mask的长度一致性在训练前会被校验len(loss_mask) ! response_length直接报错。5.3 trainable_agents.py异步多轮决策循环trainable_agents.py 中的TrainableToolCallingAgent继承 tau-bench 原生的ToolCallingAgent并混入TrainableAgentMixin把同步交互改造为基于 sglang HTTP 服务的异步循环。核心方法asolvemax_num_steps30每轮迭代做五件事构造 prompt 并发起生成用state.tokenizer.apply_chat_template(..., add_generation_promptTrue, toolsself.tools_info)渲染带工具定义的对话POST 到http://{sglang_router_ip}:{sglang_router_port}/generate_call_llm复用 slime/utils/http_utils.py 的异步post处理中止finish_reason.type abort或工具解析失败时置Status.ABORTED并提前返回解析工具调用用 OpenAI 适配器把 sglang 回复解析为normal_text calls执行环境动作call_to_action_sglang把解析结果转成 tau-bench 的Action默认动作是RESPOND_ACTION_NAME即直接回复用户随后env.step(action)获得observation、reward、done动作名非respond时以role: tool追加观察否则以role: user追加记录 token 与 loss mask每轮通过_get_token_delta计算新增 token 及其 loss mask累积到response_token_ids与loss_masks。循环结束后根据done标记判定COMPLETED或TRUNCATED_build_final_result汇总奖励、对话、loss mask 与完整 token 序列并生成response拼接所有 assistant 消息。关于工具调用指令mixin 还通过_reformulate_tool_call把 Qwen 默认的 You may call one or more functions… 提示替换为 tau-bench 友好的约束每轮至多一次工具调用、必须等待工具系统确认才能向用户确认执行等。5.4 token_delta.py多轮对话的 token 增量与 loss mask多轮场景下不能简单把整段对话全部做 loss必须精确知道本轮新增加的内容以及哪些 token 参与 RL 损失。get_token_deltatoken_delta.py按最后一条消息的角色分两种情况最后一条是 assistant分别渲染messages[:-1]prev、加上生成提示后的generation_prompt以及完整messagescurr通过前缀切片得到本轮新增文本。默认include_generation_promptFalse新增的生成提示 token 不计入只对模型自己生成的回复部分打loss_mask1若include_generation_promptTrue首轮生成时需要把生成提示也纳入则生成提示段 mask 为 0、回复段为 1最后一条是 user/tool渲染前后两段对话做前缀差得到环境新增的观察文本loss mask 全部置 0环境观察不参与策略优化。特别地代码对 Qwen3 这类新 user 消息到来时会重写历史的推理模板做了兜底当curr不以prev开头时将最后一条 user 消息独立渲染并校验是否为curr的后缀保证增量计算在模板重写历史的情况下依然正确。对应测试 test_tau_bench_token_delta.py 用自定义 tokenizer 覆盖了历史重写与边界 token 合并两类边界情况。5.5 openai_tool_adapter.py 与 sglang_tool_parser.py工具调用解析sglang_tool_parser.py 的parse_tools复用 sglang 内部的FunctionCallParser把 tau-bench 提供的工具定义构造成Tool/Function对象后调用parse_non_stream(response)返回{normal_text, calls}openai_tool_adapter.py 的OpenAICompatibleToolCallAdapter进一步把解析结果封装为 OpenAI 兼容的 assistant 消息tool_calls数组 id/function字段并提供_call_to_action_sglang把首个工具调用转成 tau-benchAction以及get_openai_tools_format输出 OpenAI 格式的工具定义便于其他 Agent 框架复用。六、训练注意事项与常见问题排查Agent 策略目前仅tool-calling已实现act/react/few-shot会触发NotImplementedError需修改 trainable_agents.py 的agent_factory自行扩展。用户模拟器 APIGEMINI_API_KEY必须配置为有效 KeyGemini 返回并发限制错误时调低--sglang-server-concurrency如 32或更换更廉价的user_model。partial rollouttau-bench 交互不支持 partial rolloutgenerate开头有显式断言。MLA 模型若改用带 MLA 的模型需注释掉--attention-backend flash。显存控制--recompute-*与--max-tokens-per-gpu需按实际 GPU 显存调整sglang_mem_fraction_static0.7是 sglang 侧的另一道显存闸门。checkpoint 路径--load/--saveslime 训练格式与--ref-loadmcore 参考权重路径需与第 2.3 节准备结果一致且首次训练前--load目录应为空或不存在。七、总结与扩展思路通过examples/tau-bench示例slime 展示了如何以极小侵入量将自定义的多轮 Agent 环境接入其 RL 训练管线只需实现一个async generate(args, sample, sampling_params) - Sample函数并通过--custom-generate-function-path挂载框架便会自动完成数据调度、sglang 服务管理、GRPO 优势估计与 Megatron 训练。在此基础上可进一步扩展更换env为airline并准备相应任务 JSONL替换agent_strategy增加多策略对比或参考 examples/tau-bench/README.md 以外的examples目录如 examples/retool、examples/search-r1、examples/strands_sglang复用同样的自定义生成函数机制构造其他 Agentic 训练场景。建议深入阅读 sglang_rollout.py 中generate_and_rm_group的分组采样逻辑--n-samples-per-prompt与动态采样过滤的配合以充分理解多轮轨迹的 reward 归一化与过滤机制。【免费下载链接】slimeslime is an LLM post-training framework for RL Scaling.项目地址: https://gitcode.com/GitHub_Trending/slime12/slime创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考