Prhub

#2005 [coding-agent-rl] Refactor coding-agent RL: turn-node TrajectoryManager + pluggable harness layer

原始 PR 作者 jingshenghang 合并时间 2026-06-17 10:08 文件变更 27 提交数 66 评论 61 代码增减 +4933 / -3287

执行摘要

用消息树 TrajectoryManager 替换分段轨迹,新增可插拔 harness 层

原有的轨迹管理基于手动分段(subagent/wipe/final),难以处理多轮对话中的 token 漂移和分支情况,导致训练数据质量受限。引入消息树模型可更精确地追踪每个 turn 的关系,并容忍 TITO 重标记偏移。同时将 sandbox 引导从示例层提升到库层,统一环境变量命名,便于扩展新的编码代理(如 Codex)。

强烈建议精读:设计中消息树模型和 token drift 处理是核心创新点;review 中的性能讨论(chunk 比较、dict 比较)和设计权衡(fork vs replace)有很高的学习价值。对于需要扩展新的 LLM agent 或改进多轮 RL 训练的项目,此 PR 提供了可复用的模式。

讨论亮点

性能优化(zhuzilin vs jingshenghang):zhuzilin 指出 _common_prefix_len 用 Python 级循环可能成为瓶颈。jingshenghang 采用 chunk 4096 分块比较,将 1M 长度场景耗时从 36ms 降至 12ms,并提供 benchmark 对比。函数重命名为 _common_prefix_len。

挂载点匹配方法(zhuzilin):zhuzilin 询问 dict == 能否匹配内容。jingshenghang 演示 dict == 递归匹配 key-value,但不敏感顺序,而 json.dumps 敏感顺序。采用 dict == 后,200 nodes 场景匹配耗时从 1198ms 降至 17.6ms。但顺序变化可能被 token drift 探测到。

锁与异常处理(zhuzilin):zhuzilin 质疑 adapter 中 lock 的必要性和 try-catch 覆盖。jingshenghang 确认无并发问题(CC 保证顺序),删除 lock;异常路径预期不会触发,删除 try-catch,允许异常直接失败。

文档风格(zhuzilin):zhuzilin 批评 trajectory_manager.py 的文档 AI 味太浓。jingshenghang 接受并大幅精简,只保留不变约定、陷阱和跨层合同。

测试命名(zhuzilin):zhuzilin 建议不要使用 "e2e" 命名测试,因为 slime 中 e2e 指完整训练。jingshenghang 将 test_trajectory_manager_e2e.py 更名为 test_trajectory_manager_branching.py。

CC title 请求过滤(zhuzilin & jingshenghang):zhuzilin 对 handle_title_generation 的魔鬼逻辑表示疑惑。jingshenghang 解释是 CC 发起的 session title 请求,应丢弃。最终通过使用 claude -p 启动规避,删除相关过滤代码。

billing header(zhuzilin & jingshenghang):zhuzilin 问是否新版 CC 添加 billing header。jingshenghang 确认并通过设置 CLAUDE_CODE_ATTRIBUTION_HEADER=0 关闭,删除代码过滤。

实现拆解

  1. 重写轨迹管理层:引入 TurnRecord 作为适配器与管理器之间的契约,TrajectoryManager 基于 per-session 消息树管理多轮对话。append_turn() 将每个 turn(prompt 消息 + 模型输出)插入树中,通过 dict 比较前缀消息进行挂载点匹配,支持 token drift 时的分支(fork)或替换(replace)操作。最终 get_trajectory() 将树线性化为带 loss mask 的 Sample 列表,确保只有最新叶子的 response 参与训练。性能优化包括使用 chunk 4096 比较和 dict 直接比较,挂载点匹配性能提升 10-100 倍。

  2. 构建可插拔 harness 层:创建 slime/agent/harness/ 包,定义 BaseHarness 抽象基类(install_cli, write_config, launch_and_wait 钩子)和 HarnessContext 数据类。实现 ClaudeCodeHarness 和 CodexHarness,共享 run_command 辅助函数(处理进程分离、超时、轨迹日志捕获)。BaseHarness.run() 提供模板方法,包括确保 agent 用户存在、写入配置、启动并等待完成。

  3. 重构适配器层并集成 TrajectoryManager:将 AnthropicAdapter 和 OpenAIAdapter 的公共逻辑抽取到 common.py 中的 BaseAdapter 类,包括 Session 管理、Reply 构建、_run_turn 流水线。每个适配器实例持有全局 TrajectoryManager 实例(per sid),record_turn 由适配器调用。删除旧的 AdapterChain、TokenSegment 等类,简化数据结构。环境变量统一为 SLIME_AGENT_* 和 ADAPTER_*,并提供旧名称的 fallback。

  4. 清理示例层并分离 SWE 逻辑:将 examples/coding_agent_rl/sandbox.py 删除,其核心引导功能(Node 安装、Claude Code 安装、agent 用户创建)迁移到 slime/agent/sandbox.py(作为公共库)和 examples/coding_agent_rl/swe.py(SWE 特定任务)。swe.py 专注数据集元数据提取、工作区准备、diff 捕获和评估,与 harness 完全解耦。

  5. 重组测试并增强 CI:将分散的测试文件合并为 tests/test_agent/ 包,新增 test_trajectory_manager_branching.py(约 1400 行,覆盖 98% 分支)、test_adapters.py、test_harness.py 和 test_agent_rollout_cpu.py。CI 新增 agent-test 任务,运行所有 agent 相关测试。测试采用 golden token+loss 字符串断言,确保训练数据生成正确。

文件 模块 状态 重要度
slime/agent/trajectory.py 轨迹管理 modified 8.94
slime/agent/harness/common.py harness 层 added 9.01
slime/agent/adapters/common.py 适配器公共 modified 8.75
examples/coding_agent_rl/sandbox.py 示例删除 removed 8.89
tests/test_agent/test_trajectory_manager_branching.py 分支测试 added 7.91

关键符号

record_turn get_trajectory run run_command _common_prefix_len flatten_content manager_finish_reason finish_session install_cli launch_and_wait get_metadata prepare_workspace evaluate

关键源码片段

slime/agent/harness/common.py core-logic

新增可插拔 harness 基类和辅助函数,是架构扩展性的关键

"""Harness-agnostic coding-agent lifecycle in a sandbox.
​
A harness is a swappable coding agent (Claude Code, Codex, ...). Each one
installs a CLI, writes its own config, and runs the agent against a prompt. The
shared parts (create the agent user, the run skeleton, the launch-detached-and-
poll transport) live here; adding a CLI-style harness means subclassing
BaseHarness and implementing install_cli, write_config and launch_and_wait.
The base knows nothing about the task: run() takes only generic fields
(workdir / session_id / adapter_url / prompt). Task-specific workspace prep and
scoring live in the example layer.
"""
​
from __future__ import annotations
​
import asyncio
import lzma
import os
import shlex
import shutil
import tempfile
import time
from abc import ABC, ABCMeta, abstractmethod
from dataclasses import dataclass
from pathlib import Path
​
from slime.agent import sandbox as _sandbox
from slime.agent.sandbox import Sandbox
from slime.utils.misc import SingletonMeta
​
​
class SingletonABCMeta(ABCMeta, SingletonMeta):
    """Combine abstract base class (ABC) with singleton behavior."""
    pass
​
​
EXIT_TIME_BUDGET_EXCEEDED = -1 # sentinel when agent run times out
​
​
@dataclass(frozen=True)
class HarnessContext:
    """Generic run context, free of any task fields.
​
    ``model_label`` is the model name the harness advertises to its CLI. The
    slime adapter ignores it and serves whatever upstream sglang has loaded.
    """
    workdir: str # workspace path inside sandbox
    session_id: str # session identifier for TrajectoryManager
    adapter_url: str # URL for reverse-connection LLM adapter
    model_label: str = "slime-actor"
​
​
class BaseHarness(ABC, metaclass=SingletonABCMeta):
    """Base lifecycle for a sandbox-resident coding agent."""
​
    name: str = "" # short identifier, set by subclass (e.g., "claude_code")
​
    @abstractmethod
    async def install_cli(self, sb: Sandbox) -> None:
        """Install the harness CLI into the sandbox.
        npm-packaged harnesses delegate to install_npm_cli."""
​
    @abstractmethod
    async def write_config(self, sb: Sandbox, ctx: HarnessContext) -> None:
        """Write any CLI config files into the sandbox."""
​
    @abstractmethod
    async def launch_and_wait(self, sb: Sandbox, ctx: HarnessContext, prompt: str, time_budget_sec: int) -> int:
        """Run the agent to completion and return its exit code.
        A non-interactive CLI builds one shell command and hands it to
        run_command. An interactive or long-running harness drives its own loop.
        """
​
    async def run(
        self,
        sb: Sandbox,
        *,
        workdir: str,
        session_id: str,
        adapter_url: str,
        time_budget_sec: int,
        prompt: str,
    ) -> int:
        """Template method: ensure agent user -> write config -> launch & wait.
        Workspace prep (writing problem statement etc.) is the caller's job.
        """
        await _sandbox.ensure_agent_user(sb, workdir)
        ctx = HarnessContext(
            workdir=workdir,
            session_id=session_id,
            adapter_url=adapter_url,
        )
        await self.write_config(sb, ctx)
        return await self.launch_and_wait(sb, ctx, prompt, time_budget_sec)

评论区精华

LCP 性能优化 性能

zhuzilin 担心 _common_prefix_len 的循环成为性能瓶颈。jingshenghang 采用 chunk 4096 分块比较,将耗时从 36ms 降至 12ms,并提供 benchmark 对比。

结论:采用 chunked 比较,函数重命名为 _common_prefix_len。 · 已解决

挂载点匹配方法 性能

zhuzilin 询问 dict == 能否匹配内容。jingshenghang 演示 dict == 递归匹配 key-value,但顺序不敏感,而 json.dumps 敏感顺序。采用 dict == 后匹配耗时从 1198ms 降至 17.6ms。顺序变化由 token drift 探测处理。

结论:采用 dict == 直接比较,提升性能。 · 已解决

lock 和 try-catch 必要性 正确性

zhuzilin 质疑 adapter 中 lock 的必要性和 try-catch 覆盖。jingshenghang 确认无并发问题(CC 保证顺序),删除 lock;异常路径预期不会触发,删除 try-catch,允许异常直接失败。

结论:删除 lock 和 try-catch。 · 已解决

文档 AI 味太重 style

zhuzilin 批评 trajectory_manager.py 的文档太 AI 味。jingshenghang 接受并大幅精简,只保留关键约定和陷阱。

结论:精简文档。 · 已解决

测试命名 e2e 测试

zhuzilin 建议不要使用 e2e 命名测试,因为 slime 中 e2e 指完整训练。jingshenghang 将测试文件更名为 test_trajectory_manager_branching.py。

结论:重命名测试文件。 · 已解决

CC title 请求过滤 设计

zhuzilin 对 handle_title_generation 的魔鬼逻辑表示疑惑。jingshenghang 解释是 CC 发起的 session title 请求,应丢弃。最终通过使用 claude -p 启动规避,删除相关过滤代码。

结论:删除过滤代码,改用 claude -p 启动。 · 已解决

billing header 处理 设计

zhuzilin 问是否新版 CC 添加 billing header。jingshenghang 确认并通过设置 CLAUDE_CODE_ATTRIBUTION_HEADER=0 关闭,删除代码过滤。

结论:通过环境变量关闭 billing header,删除代码过滤。 · 已解决

风险与影响

  • 训练数据生成路径改变:TrajectoryManager 的 token drift 处理逻辑(fork/replace)可能对模型训练产生不可预期的影响。但通过 98% 测试覆盖和 review 中的反复打磨,风险可控。
  • 环境变量迁移:旧环境变量(SWE_HOST_*, SLIME_HEAD_HOST, SHIM_*)被替换为 SLIME_AGENT_* 和 ADAPTER_*。虽然通过 _getenv 提供向后兼容,但用户仍需更新部署脚本。文档已同步更新。
  • harness 单例状态:BaseHarness 采用 SingletonABCMeta,若安装或配置有残留状态可能影响后续调用。但 harness 设计为无状态(仅在运行时安装 CLI、写入配置),风险低。
  • 适配器重构:旧适配器 API(如 AdapterChain)被移除,外部代码若直接使用需更新。但示例已更新反映新 API。
  • 对用户:使用 --custom-generate-function-path examples.coding_agent_rl.generate.generate 的用户需要同步更新其环境变量和可能的自定义适配器。但 PR 提供了向后兼容层。
  • 对系统:新轨迹系统使用消息树,内存开销略增(树结构 vs 线性段),但性能优化抵消。训练数据生成质量提升,多轮对话中的 token 漂移被容忍,减少数据浪费。
  • 对团队:模块化明显改善,harness 和 adapter 分离,便于添加新编码代理(如 Codex)。测试覆盖从低到高,CI 增加 agent 专项测试,降低回归风险。文档更新(README)简化部署说明。
核心路径变更 环境变量迁移 训练数据生成路径改变 适配器 API 破坏

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论