TRL documentation

用于在环境中训练 LLM 的 OpenEnv 集成

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

用于在环境中训练 LLM 的 OpenEnv 集成

OpenEnv 是一个用于定义、部署和交互强化学习(RL)及智能体工作流中环境的开源框架。它提供了标准化的环境交互 API,并支持将环境作为后端服务器运行(通过 WebSocket 或容器化执行)。您可以在 Hugging Face Hub 上找到一系列可直接使用的 OpenEnv 环境。

本指南介绍了如何将 OpenEnv 与 TRL 集成。有关 OpenEnv 本身的更多信息,请参阅OpenEnv 文档。

现成的 OpenEnv 示例:grpo_echo(最简示例)、grpo_catch、grpo_wordle、grpo_sudoku、grpo_multi_env、grpo_browsergym、grpo_carla,以及 async_grpo_harbor(本地和 HF Jobs)。

何时使用环境

GRPOTrainer 可用于训练智能体。对于智能体任务,它支持两种模式:tools,模型可以调用外部函数,但每次调用都是无状态且相互独立的;以及 environments,它能在多轮交互中保持状态,从而实现真正的多轮交互,其中智能体的行为会影响后续的观测结果。当连续性至关重要时,请使用 environments,例如在游戏导航、浏览网页,或任何智能体下一步看到的内容取决于其先前行为的任务中。

选择环境集成

OpenEnv 是本文档中介绍的默认路径。另外两个集成——OpenReward 和 Harbor——遵循相同的 environment_factory 契约,在 TRL 层面可以互换使用。请参阅 GRPO 指南中的环境集成对比,选择与您的任务生态相匹配的集成。

安装

OpenEnv 环境以 Hugging Face Spaces 的形式托管,它们同时也是可通过 pip 安装的 Git 仓库:

# Echo environment
pip install "openenv-echo-env @ git+https://huggingface.co/spaces/openenv/echo_env"

# Wordle (TextArena) environment
pip install "openenv-textarena @ git+https://huggingface.co/spaces/openenv/wordle"

# Catch (OpenSpiel) environment
pip install "openenv-openspiel-env @ git+https://huggingface.co/spaces/openenv/openspiel_env"

这将安装环境客户端(例如 EchoEnv),它通过 WebSocket 与远程环境服务器进行通信,同时安装动作/观测模型以及所有必需的依赖项(包括 openenv)。

您可以在任何环境的 HF Space 页面上找到安装命令。点击 ⋮(三个点) 菜单并选择 “使用此 Space” 即可查看安装说明。

您也可以通过 pip install "openenv>=0.3.1" 从 PyPI 安装核心包,但请注意,特定环境的依赖项可能需要单独安装。

用于开发,您可以克隆 OpenEnv 仓库并在本地安装:

git clone https://github.com/openenv/OpenEnv.git
cd OpenEnv/envs/echo_env
pip install -e .

TRL 中的每个环境脚本都包含内联依赖元数据(PEP 723),因此您也可以直接使用 uv 运行它们:

uv run examples/grpo_echo/grpo_echo.py

这会自动在隔离的虚拟环境中安装所需的环境包。

快速入门

理解集成的最快方式是查看一个完整的示例。echo.py 脚本使用 Echo 环境 来训练模型,该环境根据补全文本的文本长度给予奖励:

from datasets import Dataset
from echo_env import EchoEnv
from echo_env.models import EchoAction

from trl import GRPOConfig, GRPOTrainer

ENV_URL = "https://huggingface.co/proxy/openenv-echo-env.hf.space"

class EchoToolEnv:
    def __init__(self):
        self.env = EchoEnv(base_url=ENV_URL)
        self.reward = 0.0

    def reset(self, **kwargs) -> str | None:
        self.reward = 0.0
        return None

    def echo(self, message: str) -> str:
        """
        Echo the message back from the environment.

        Args:
            message: The message to echo

        Returns:
            The echoed message.
        """
        observation = self.env.step(EchoAction(message=message))
        self.reward = observation.observation.reward
        return observation.observation.echoed_message

def reward_func(environments, **kwargs):
    return [env.reward for env in environments]

dataset = Dataset.from_dict(
    {"prompt": [[{"role": "user", "content": "Try to echo 'Hello World!' in the environment."}]] * 64}
)

trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=dataset,
    reward_funcs=reward_func,
    args=GRPOConfig(
        chat_template_kwargs={"enable_thinking": False},
        log_completions=True,
    ),
    environment_factory=EchoToolEnv,
)
trainer.train()

就是这样。以下是其内部工作原理:

  1. environment_factory=EchoToolEnv:训练器会为每次生成创建一个 EchoToolEnv 实例(传入类,而非实例)。
  2. reset() 在每个回合开始时被调用以初始化状态。返回一个观测字符串(或 None)。
  3. 工具发现:训练器会发现环境实例上的所有公共方法(此处为 echo()),并将它们暴露为函数调用工具。每个方法都必须具有包含类型化参数的规范文档字符串,训练器将使用它来构建工具架构。
  4. 多轮循环:训练器生成补全内容,解析工具调用,执行 echo(),将结果追加到上下文中,然后再次生成,直到模型停止调用工具或达到 max_completion_length 为止。
  5. 奖励函数:在回合结束后(环境重置前)从每个环境实例中读取 env.reward。
# Run the example
python examples/grpo_echo/grpo_echo.py

# Customize model and environment URL
python examples/grpo_echo/grpo_echo.py --model Qwen/Qwen3-0.6B --env-host https://huggingface.co/proxy/openenv-echo-env.hf.space

以下是训练过程中的奖励曲线:

您可以在本页顶部的笔记中探索更多现成的 OpenEnv 示例。

environment_factory 的工作原理

TRL 的 GRPOTrainer 通过 environment_factory 参数支持交互式环境训练。提供该参数后,训练器会自动处理多轮工具调用循环:生成补全、解析工具调用、在环境中执行它们,并将结果反馈给模型。全程无需自定义 rollout 代码。

Environment 类要求

您的环境类必须遵循以下规则:

  • __init__(self) (可选):如果提供,则不能接受任何参数。使用它来初始化状态或客户端。如果需要外部配置(例如 URL),请从外围作用域或模块级变量中获取。
  • reset(self, **kwargs):在每个 episode 开始时调用。以关键字参数的形式接收所有数据集列。返回一个字符串形式的观测值(若没有初始观测值,则返回 None)。
  • 工具方法:除 reset 外,任何不以 _ 开头的公开方法都会自动暴露为工具。每个工具方法都必须包含带有 Args: 描述的文档字符串,因为训练器会使用这些描述来为模型生成工具架构。

环境类提示

  • 奖励状态:你可以在环境实例上存储任意状态(例如 self.reward、self.done 等),并通过 environments 参数在奖励函数中访问它们。有关此模式的示例,请参阅 快速入门指南。
  • 错误处理:如果工具方法抛出异常(例如 ValueError("Game over.")),训练器会捕获该异常,并将错误消息作为工具响应反馈给模型。这是指示操作无效或回合已结束的建议方式。
ENV_URL = "https://huggingface.co/proxy/my-env.hf.space"

class MyEnv:
    def __init__(self):
        self.client = MyClient(base_url=ENV_URL)  # captured from enclosing scope
        self.reward = 0.0

    def reset(self, **kwargs) -> str | None:
        self.reward = 0.0
        return "Initial observation for the model"

    def my_tool(self, arg1: str, arg2: int) -> str:
        """
        Description of what this tool does.

        Args:
            arg1: Description of arg1
            arg2: Description of arg2

        Returns:
            The result message.
        """
        self.reward = 1.0
        return "Tool result"

工具必须是独立的方法,并配有描述性的名称和带类型的参数(例如 guess(word: str)、move(direction: str))。我们不建议使用像 step(action) 这样的通用方法,因为模型需要具有意义的工具名称和参数描述才能学会调用工具。

奖励函数

奖励函数会接收 environments 参数(一个环境实例列表),因此您可以访问回合期间存储的任何状态:

def reward_func(environments, **kwargs) -> list[float]:
    return [env.reward for env in environments]

有关奖励函数的更多信息,请参阅 GRPO - 自定义奖励函数。

奖励函数技巧

在使用 OpenEnv 环境和 GRPO 时,我们发现以下几点很有帮助:

  • 确保组内存在差异。GRPO 在组内比较生成结果,因此相对排名比绝对数值更重要,而所有生成结果获得相同奖励的组不会提供任何信号。当模型有时能成功时,二元奖励(成功为 1.0,否则为 0.0)是清晰的。当模型很少成功时,请添加部分得分:数独示例将已填充单元格的分数加到解题奖励中,Wordle 示例则加入最佳猜测中的绿色字母数。
  • 检查最终状态,而非路径。 在可能的情况下,让环境来评判结果(例如,“模型是否解开了谜题?”),而不是检查它是否遵循了特定的操作序列。这能让模型自由地发现其自身的策略。
  • 在训练前测试你的奖励。 手动运行几个回合(参见 Wordle 示例笔记本),以确认环境返回的奖励合理。如果一个能力较强的模型无法取得高于随机基线的分数,则可能需要调整奖励信号。

多轮对话中的 max_completion_length

max_completion_length 参数限制的是整个多轮对话中所有 token 的总数(包括所有模型生成内容和工具结果的总和),而不仅仅是单次生成。对于轮次较多的环境(例如需要数十步操作的数独),您可能需要增大该值:

args = GRPOConfig(
    max_completion_length=4096,  # default is usually 256-1024, increase for long episodes
    # ...
)

如果回合被提前中断(模型在游戏中途停止),这很可能就是原因。

针对 Space 的长时间运行

Space 可能会关闭长时间保持打开的 WebSocket 连接(我们观察到大约每 2 小时发生一次)。在 0.11.0 及更早版本的 OpenEnv 客户端中,随后会重新连接到新的环境会话,如果当时有游戏正在进行,下一次 step() 将在客户端从未重置过的游戏上运行,且不会报错。从 OpenEnv 0.12.0 开始,客户端会抛出 SessionLostError,从而导致游戏失败。为了避免训练过程中出现断连,请在 reset() 中打开一个新的客户端,使每个连接仅持续一个 episode,并将 MAX_CONCURRENT_ENVS 设置为略高于同时运行的游戏数量,因为旧会话是异步释放的:

def reset(self, **kwargs) -> str | None:
    self.client.close()
    self.client = TextArenaEnv(base_url=env_url).sync()
    result = self.client.reset()
    ...

高级示例:Wordle

让我们使用 TextArena 环境训练一个模型来玩 Wordle。这展示了多轮交互、累积反馈处理、拆分为多个项的奖励,以及通过异常终止回合。

OpenEnv Wordle 教程涵盖了整个流程:SFT 数据、SFT 预热、本 GRPO 步骤以及评估,并附有硬件和成本说明。此示例的 notebook 版本为 grpo_wordle.ipynb。

TextArena 环境

TextArena 是一个开源的竞技类文本游戏合集,旨在通过 Wordle、贪吃蛇、井字棋等文本游戏来评估大语言模型的推理能力。

image of TextArena

为什么选择 Wordle?

Wordle 是环境强化学习的一个良好基准,因为它需要对反馈进行推理,完全基于文本,且参数量仅为 1B 的模型也能在其中取得提升。

Wordle 是如何运作的? Wordle 是一款猜词游戏,玩家需要在 6 次尝试内猜出一个 5 个字母的单词。每次猜测后,环境会提供逐字母的反馈:

G U E S S
X G Y X X

X = 不在单词中,G = 位置正确(绿色),Y = 位置错误(黄色)。在此,“U” 正确且位置正确,“E” 在单词中但位置错误。

从 SFT 预热开始

Qwen3-1.7B 几乎从不赢得游戏,因此从基础模型出发的 GRPO 几乎没有可学习的内容:一组没有一局获胜的游戏无法提供任何信号。该示例从 openenv/Qwen3-1.7B-Wordle-SFT 开始,该模型使用 求解器在同一工具调用格式下进行的 1,000 局游戏 训练了一个 epoch,并在每次猜测前加入简短的推理。

Environment 类

WordleEnv 类封装了 TextArena 客户端,将 guess() 暴露为工具,并保留了奖励计算所需的内容:

import re

from textarena_env import TextArenaAction, TextArenaEnv

FEEDBACK = re.compile(r"Feedback:\n([A-Z ]+)\n([GYX ]+)")

class WordleEnv:
    def __init__(self):
        self.client = TextArenaEnv(base_url="https://<user>-wordle.hf.space")

    def reset(self, **kwargs) -> str | None:
        result = self.client.reset()
        self._last_full_feedback = result.observation.messages[0].content
        self.done = False
        self.won = False
        self.num_guesses = 0
        self.best_greens = 0
        self.invalid_moves = 0
        return self._last_full_feedback

    def guess(self, guess: str) -> str:
        """
        Make a guess in the Wordle environment.

        Args:
            guess: The guessed word, formatted as '[abcde]'

        Returns:
            The feedback message from the environment.
        """
        if self.done:
            raise ValueError("Game over.")
        result = self.client.step(TextArenaAction(message=guess))
        _full_feedback = result.observation.messages[0].content
        feedback = _full_feedback[len(self._last_full_feedback):]
        self._last_full_feedback = _full_feedback
        if "invalid move" in feedback.lower():
            self.invalid_moves += 1
        else:
            self.num_guesses += 1
            match = FEEDBACK.search(feedback)
            if match:
                self.best_greens = max(self.best_greens, match.group(2).count("G"))
            self.won = result.done and result.reward >= 1.0
        self.done = result.done
        return feedback

关键设计选择:

  • reset() 返回初始游戏消息,作为模型看到的第一条观察。
  • guess() 是唯一的工具。模型每轮都会调用它并传入一个 5 个字母的单词。
  • 累积反馈切片:TextArena 每回合都会返回完整的游戏历史。我们只切分出新增部分,以避免重复上下文。
  • done 时抛出异常:如果模型在游戏结束后仍尝试猜测,guess() 会抛出 ValueError。训练器会捕获该异常,并将 "Game over." 作为工具响应反馈给模型,同时在 tools/failure_frequency 中计入一次。

奖励函数与训练

环境自身的奖励在猜对一次或六次后支付的获胜奖励相同,并且仅在游戏结束时给予部分分数。在开启思考功能的情况下,策略可能会学会减少思考,并利用全部六次猜测机会来结束游戏。以下三个奖励项可以避免这种情况:

from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer

def reward_win(environments, **kwargs) -> list[float]:
    return [1.0 + 0.5 * (1 - env.num_guesses / 6) if env.won else 0.0 for env in environments]

def reward_partial(environments, **kwargs) -> list[float]:
    return [0.0 if env.won else 0.1 * env.best_greens for env in environments]

def reward_invalid(environments, **kwargs) -> list[float]:
    return [-0.1 * env.invalid_moves for env in environments]

prompt = """You are an expert Wordle solver with deep knowledge of English vocabulary...
Use the tool `guess` to make a guess."""

dataset = Dataset.from_dict({"prompt": [[{"role": "user", "content": prompt}]] * 600})

trainer = GRPOTrainer(
    model="openenv/Qwen3-1.7B-Wordle-SFT",
    reward_funcs=[reward_win, reward_partial, reward_invalid],
    train_dataset=dataset,
    args=GRPOConfig(
        use_vllm=True,
        vllm_mode="colocate",
        chat_template_kwargs={"enable_thinking": True},
        max_completion_length=4096,
        max_tool_calling_iterations=12,
        vllm_importance_sampling_mode="token_truncate",
        num_generations=16,
        per_device_train_batch_size=8,
        gradient_accumulation_steps=8,
        learning_rate=1e-6,
        lr_scheduler_type="constant",
        max_steps=150,
    ),
    environment_factory=WordleEnv,
)
trainer.train()

获胜的奖励随获胜速度加快而增加(一次猜中得 1.42,六次猜中得 1.0),无论游戏是否结束,失败都会根据该局最佳猜测给予奖励,每次无效操作扣除 0.1。TRL 会分别记录每一项,因此 rewards/reward_win/mean 会随着胜率上升而增加。

运行示例

每个步骤会进行 64 局游戏,每局对应一个环境会话。复制 openenv/wordle Space,并在您的副本中设置 MAX_CONCURRENT_ENVS=64 变量。

colocate
server

Colocate 模式(1 个 GPU,推荐)

python examples/grpo_wordle/grpo_wordle.py --vllm-mode colocate --env-url https://<user>-wordle.hf.space

这会在与训练相同的进程中运行 vLLM,仅需一块 GPU。在我们的运行中,150 步在单块 A100 80GB 上耗时约 2 小时 15 分钟。

结果

在我们的运行中,每个模型进行一次训练,在相同的 30 个单词上以贪心方式生成:

模型胜利每次获胜的猜测次数无效的猜测
Qwen3-1.7B0/30-24
+ SFT10/303.241
+ SFT + GRPO(模型)20/304.13

GRPO 主要消除了重复和捏造的词,并赢得了更多游戏。大多数新胜局需要 4 到 6 次猜测,因此每次获胜所需的猜测次数有所增加。这 30 个词来自 TextArena 用于生成秘密词的 156 个名词,这些名词在训练过程中也被模型见过,而一次新的运行结果可能会相差几个胜局(使用相同配方的两次 SFT 运行分别赢得了 10 局和 13 局)。在 Wordle Arena 中观看这三个模型的对局,并在 Trackio 上查看训练曲线。

多环境训练

你可以同时在多个环境中训练单个模型。当你希望模型并行学习不同技能时,这非常有用。例如,在同一次训练运行中同时玩 Wordle(语言推理)和 Catch(空间推理)。

核心思想是创建一个元环境类,该类封装多个环境,并利用数据集列将每个样本路由到正确的环境中。

工作原理

  1. 在数据集中添加一个 "env" 列(或类似的列),用于标识每个样本所属的环境。
  2. 在 reset(**kwargs) 中,读取 kwargs["env"] 以选择该回合的活跃环境。
  3. 暴露来自所有环境的工具;训练器会发现所有公共方法。
  4. 为每个环境使用独立的奖励函数,对于不属于该环境的样本返回 None。TRL 会使用 nansum/nanmean 来处理 None 值。

示例:Wordle + Catch

grpo_multi_env.py 脚本会同时在 Wordle 和 Catch 上进行训练:

class MultiEnv:
    def __init__(self):
        self._wordle_client = None
        self._catch_client = None
        self.active = None
        self.reward = 0.0
        self.done = False

    def reset(self, **kwargs) -> str | None:
        self.active = kwargs.get("env", "wordle")
        self.reward = 0.0
        self.done = False

        if self.active == "wordle":
            if self._wordle_client is not None:
                try:
                    self._wordle_client.close()
                except Exception:
                    pass
            self._wordle_client = TextArenaEnv(base_url=WORDLE_URL)
            result = self._wordle_client.reset()
            self._last_full_feedback = result.observation.messages[0].content
            self.reward = 0.0
            return self._last_full_feedback
        elif self.active == "catch":
            if self._catch_client is not None:
                try:
                    self._catch_client.close()
                except Exception:
                    pass
            self._catch_client = OpenSpielEnv(base_url=CATCH_URL)
            result = self._catch_client.reset()
            self.done = result.observation.done
            return _format_catch_obs(result.observation.info_state)

    # Wordle tool
    def guess(self, guess: str) -> str:
        """Make a guess in the Wordle environment. ..."""
        ...

    # Catch tools
    def move(self, direction: str) -> str:
        """Move the paddle left or right. ..."""
        ...

    def stay(self) -> str:
        """Do nothing and let the ball fall one step. ..."""
        ...

关键模式:

  • 延迟客户端初始化:在 reset() 中创建客户端,而不是在 __init__() 中,以避免不必要的 WebSocket 连接。
  • 关闭后重新打开:在创建新客户端之前,先关闭上一个客户端,以避免服务器容量错误。
  • kwargs 路由:数据集的 "env" 列会作为关键字参数传递给 reset()。
  • 所有工具同时暴露:无论当前处于何种环境,模型都会将 guess、move 和 stay 视为可用工具。如果它调用了错误的工具(例如在 Wordle 中调用 move),该方法会引发一个 ValueError,训练器会妥善捕获该异常。在实际操作中,模型会根据系统提示词学会使用正确的工具。

每个环境的奖励函数

每个奖励函数对其他环境中的样本返回 None:

def wordle_reward(environments, **kwargs) -> list[float | None]:
    return [env.reward if env.active == "wordle" else None for env in environments]

def catch_reward(environments, **kwargs) -> list[float | None]:
    rewards = []
    for env in environments:
        if env.active != "catch":
            rewards.append(None)
        elif env.done:
            rewards.append(max(env.reward, 0.0))
        else:
            rewards.append(0.0)
    return rewards

TRL 内部会将 None 转换为 nan,并使用 nansum/nanmean 进行聚合,因此每个样本仅由其相关的奖励函数进行评分。

带有环境路由的数据集

n = 500
dataset = Dataset.from_dict({
    "prompt": (
        [[{"role": "user", "content": wordle_prompt}]] * n
        + [[{"role": "user", "content": catch_prompt}]] * n
    ),
    "env": ["wordle"] * n + ["catch"] * n,
})

运行多环境示例

python examples/grpo_multi_env/grpo_multi_env.py \
    --wordle-url https://huggingface.co/proxy/openenv-wordle.hf.space \
    --catch-url https://huggingface.co/proxy/openenv-openspiel-env.hf.space \
    --vllm-mode colocate \
    --gradient-accumulation-steps 4 \
    --num-generations 8

在跨多个环境进行训练时,应监控每个奖励函数的指标(train/reward_func_0、train/reward_func_1 等),而不是合并后的 train/reward。合并指标会在不同环境之间交替变化,因此可能显得噪声较大。

运行环境

使用 environment_factory 时,训练器会自动连接到环境服务器。您只需确保服务器正在运行即可。运行 OpenEnv 环境服务器有三种方式:

space
docker
local

连接到远程 Hugging Face Space (最简单)

大多数示例脚本默认使用托管的 Space(无需任何配置):

env = EchoEnv(base_url="https://huggingface.co/proxy/openenv-echo-env.hf.space")

对于训练,请将 Space 复制到您自己的账户,以避免并发问题。训练器会打开 N 个并发的 WebSocket 连接(每个生成任务一个),而共享的 Space 可能不支持此操作。有关详细信息,请参阅服务器并发。

环境目录

探索当前受维护环境目录的最佳方式是访问 OpenEnv 官方目录。

要创建自己的环境,请参阅使用 OpenEnv 构建自己的环境指南。环境与 Hub 紧密集成,因此你可以推送新环境供社区复用。

服务器并发

使用 environment_factory 时,trainer 会创建 N 个环境实例(每次生成对应一个),每个实例都会向服务器打开一个 WebSocket 连接。默认情况下,OpenEnv 服务器仅允许 1 个并发会话,这会导致训练过程中出现失败。

要支持并行训练,请为服务器配置并发:

  1. 在您的环境文件中,声明对并发会话的支持:
SUPPORTS_CONCURRENT_SESSIONS: bool = True
  1. 在您的服务器应用中,设置并发限制:
app = create_app(
    create_my_environment,
    MyAction,
    MyObservation,
    max_concurrent_envs=64,  # match or exceed generation_batch_size
)

max_concurrent_envs 应 ≥ generation_batch_size(其默认值为 per_device_train_batch_size × gradient_accumulation_steps)。例如,当 gradient_accumulation_steps=64 且 batch size 为 1 时,至少需要 64 个并发会话。

environment_factory 与 rollout_func

GRPOTrainer 支持两种基于环境的训练方法:

  • environment_factory(推荐):您定义一个包含工具方法的环境类,训练器会自动处理生成、工具调用解析以及多轮循环。本指南全程采用此方法。
  • rollout_func:您自行编写完整的生成与环境交互循环。这使您可以完全控制补全的生成方式、工具的执行方式以及奖励的计算方式。

当 environment_factory 不符合您的使用场景时,请使用 rollout_func。例如,外部智能体服务器,其中外部服务器拥有生成循环并管理其自身的智能体-环境交互协议。

在测试框架上训练:训练真实的编码智能体(Harbor)

上述集成属于白盒:TRL 自行驱动多轮循环。它会采样每一轮,解析工具调用,执行这些调用,并将结果反馈回去。

某些智能体无法以这种方式驱动,因为它们拥有自己的循环。像 opencode 或 claude-code 这样的生产级编码智能体框架拥有自己的规划器、工具集、上下文管理和停止条件。你希望训练的是该智能体本身,而不是它的重新实现。

为此,TRL 提供了一个基于 AsyncGRPOTrainer 构建的实验性 黑盒(循环拥有者) 路径:HarnessAsyncGRPOTrainer 驱动一个 OpenEnv 会话工厂。智能体及其沙箱来自 Harbor,并通过 openenv harbor serve 经由 OpenEnv 提供服务,因此 Harbor 安装的每个智能体都可以以相同的方式进行训练:您只需使用 --harness 进行选择。请参阅 examples/async_grpo_harbor 获取完整示例,其中包括一个 Hugging Face Jobs 启动器。

此路径需要 openenv>=0.7.0,这是首个包含 TrainingTrace 捕获 API 的版本:pip install "openenv>=0.7.0"。

工作原理

TRL 在此处不对每一轮进行采样。智能体自行运行至完成,然后 TRL 读取其执行结果:

  1. OpenEnv Harbor 服务器会在沙箱(例如 E2B)中针对 Harbor 任务启动智能体。智能体发出的每次模型调用都会通过服务器的捕获代理转发至您的 vLLM 服务器,该服务器会记录引擎自身的提示词和补全 token id、采样的 logprobs 以及逐 token 的损失掩码。
  2. 当智能体停止时,任务自身的验证器会对结果进行评分,TRL 从 OpenEnv 接收经过验证的 TrainingTrace,并根据其 token id 和掩码构建训练行。TRL 不会重新对捕获的提示进行分词。
  3. GRPO 在这些行上进行训练。奖励通过组相对优势传播到每个受训练的 token,新的权重随后同步回同一个 vLLM 服务器。

每个 rollout 都在独立的会话中运行,拥有各自的沙箱和各自的捕获会话。这种隔离至关重要,原因有二:代理必须精确捕获该 rollout 的 token,且一个 rollout 不得干扰另一个 rollout。max_concurrent_rollouts 设置控制同时运行的 rollout 数量,而服务器的 MAX_CONCURRENT_ENVS 必须至少为 max_concurrent_rollouts + 1(工厂会保留一个额外连接用于任务元数据)。

接线

HarnessAsyncGRPOTrainer 与 AsyncGRPOTrainer 一样接收模型、配置和数据集,但使用会话工厂来替代 reward_funcs、tools 和 environment_factory。如果没有 harness_adapter,智能体将自行管理其循环。训练器会使用自身的采样策略来构建会话工厂,因此您应传入一个可调用对象(例如 partial(HarborSessionFactory, ...)),而不是一个已构建好的工厂。每个会话都会返回一个 OpenEnv TrainingTrace,其中包含精确的引擎 token、行为 logprobs 以及预填充掩码。rollout_reward_fn 是可选的;如果未提供,TRL 将使用会话的验证器奖励。

from functools import partial

from harbor_env.harness import HarborSessionFactory

from trl import AsyncGRPOConfig
from trl.experimental.async_grpo.openenv_harness import HarnessAsyncGRPOTrainer

factory = partial(
    HarborSessionFactory,
    "http://localhost:8200",       # the OpenEnv Harbor server
    split="<hf-dataset>",          # a Harbor task dataset the server serves
    harness="mini-swe-agent",      # or "opencode", "claude-code", ...
    sandbox="e2b",
    llm_url=vllm_url,
    model=model,
)
config = AsyncGRPOConfig(num_generations=8, max_concurrent_rollouts=8, vllm_server_base_url=vllm_url)
trainer = HarnessAsyncGRPOTrainer(
    model=model,
    harness_session_factory=factory,  # called by the trainer with its sampling policy
    args=config,
    train_dataset=dataset,
    rollout_reward_fn=my_reward,      # outcome -> float | None
)
trainer.train()

您需要定义的内容

OpenEnv 负责捕获验证和 token 选择。TRL 将生产者(producer)的完整序列掩码转换为其内部的补全掩码,打包生成的行并应用 GRPO。它不会重建 token,也不会推断哪些调用属于智能体(agent)。

rollout_reward_fn:如何对 rollout 进行评分

Callable[[HarnessRolloutOutcome], float | None]

TRL 并不了解您的任务中“成功”的含义,因此您需要将每个完成的 rollout 转换为标量奖励。该函数接收一个 HarnessRolloutOutcome,用于描述智能体执行的操作:

  • env_reward (float | None):来自任务验证器的奖励,如果 rollout 无法评分,则为 None。
  • completion (list[dict]): 最终的消息记录。
  • trace (list[TraceEntry]):所选的 OpenEnv 捕获记录,包括零掩码轮次。
  • tool_call_count 和 tool_failure_count(int):智能体发起的工具调用次数,以及其中看似失败的次数。
  • tool_calls_by_name (dict[str, int]): 每个工具的调用次数,例如 {"bash": 3, "edit": 2}。
  • timed_out (bool): 智能体是否用完了其时间预算。

返回一个 float,或返回 None 以标记该 rollout 无法评分,从而将其从组基线中剔除,而不是将其计为零。如果未传入此函数,则直接使用原始的 env_reward。Harbor 示例将验证器的正确性作为主要信号,仅在任务解决时添加一个小的工具效率奖励。

为什么不直接使用验证器? 验证器是环境的一部分,它只回答一个问题:“结果有多正确”,这使其保持简洁,并可复用于评估。用于训练的奖励是一个独立的、训练时的决策(例如将分数二值化、惩罚退化行为、丢弃无法打分的 rollout)。此外,它还需要验证器从未接触过的信号,因为验证器只检查最终结果,而 rollout_reward_fn 还能获取轨迹(工具调用次数、timed_out、trace)。例如,某个 rollout 可能通过了部分测试,但从未运行 bash;只有 rollout_reward_fn 能看到这一点并对其进行惩罚。

Token 的资格来自 OpenEnv

session.fetch_training_trace() 返回一个经过验证的 TrainingTrace。每一轮对话都具有唯一的 node_id、prompt_token_ids、completion_token_ids、per_token_logps,以及覆盖 prompt 和 completion 的必需 loss_mask。Prompt 位置的掩码值为零。部分 completion 掩码会被保留,包括 assistant 响应内部的零值。

例如,提示 [10, 11]、补全 [20, 21, 22] 和掩码 [0, 0, 1, 0, 1] 会在标记 20 和 22 上进行训练。标记 21 保留为上下文。完全被掩码的补全会保留用于上下文和用量统计。

生产者会排除辅助调用和被丢弃的重试。此 API 中没有 train_turn_fn 或 agent_turn_fn 参数。现有的集成必须将选择逻辑移至其 OpenEnv 掩码中。任务的验证器独立于 token 的资格。

格式错误的捕获会在 OpenEnv 边界处失败并停止训练。传输失败仍无法评分。重写后的历史会开启新的行,同时保留采样的 token 和 logprobs。多行仍然共享 rollout 奖励;此 API 不会改变 GRPO 的加权方式。

指标

Harness 训练器以相同的名称记录 AsyncGRPOTrainer 的 工具指标,区别如下:

  • 当智能体自行管理其循环(没有 harness_adapter)时,失败情况是根据最后一次请求中的工具结果来估算的:如果结果的文本中包含 error、failed、traceback 或 exception(不区分大小写),则该结果被视为失败;具有相同工具且前 200 个字符相同的结果仅计算一次。tools/latency_s 和 tools/<name>/latency_s 不会被记录,因为 OpenEnv 的跟踪信息中没有时间戳。
  • 使用 harness_adapter 时,如果 harness 报告了某个工具调用的错误,则该工具调用被视为失败。一次模型回复与下一次请求之间的时间会平均分配给该轮次的所有工具调用,并记录为它们的 tools/latency_s,因此其中包含了 harness 自身的开销。最后一轮不进行计时。
  • tools/unknown_name_count、tools/parallel_calls_mean 和 tools/loop_exhausted_frac 来自内置的工具循环,不会被记录。
  • Trainer 本身不包含任何工具,因此 tools/call_frequency 和 tools/failure_frequency 仅统计至少有一个 rollout 调用了工具的组数。
  • rollout/correctness_mean 仅在提供 rollout_reward_fn 时记录:即验证器的 env_reward 在所有具有该奖励的 rollout 上的平均值。

要求

OpenEnv Harbor 服务器负责管理数据集和沙箱。它需要沙箱后端的凭据(例如 E2B_API_KEY):

MAX_CONCURRENT_ENVS=9 openenv harbor serve --dataset <hf-dataset> --port 8200 --capture-port 8300 --expose gradio

vLLM 服务器必须暴露工具调用和真实的 token id,并启用 NCCL 权重同步,以确保 agent 始终命中当前策略。Rollouts 和权重更新必须使用同一个 vLLM 服务器:

vllm serve <model> \
    --enable-auto-tool-choice --tool-call-parser <parser> \
    --logprobs-mode processed_logprobs \
    --return-tokens-as-token-ids \
    --weight-transfer-config '{"backend":"nccl"}'

Loop-owning 训练位于 trl.experimental 下,其 API 可能会发生变化。该示例固定了经过测试的 OpenEnv 版本;有关确切配置,请参阅其 README。

Update on GitHub