TRL documentation

環境を用いたLLMの学習のためのOpenReward統合

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.17.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

環境を用いたLLMの学習のためのOpenReward統合

OpenRewardは、Open Reward Standard (ORS)を基盤とするRL環境のためのオープンエコシステムです。ORSは、環境がタスク、ツール、セッション、報酬をどのように公開するかを定義する、公開された言語非依存のHTTP/SSEプロトコルです。ORSは単なるプロトコルであるため、同じ環境をOpenRewardプラットフォーム上で実行したり、任意のコンテナサービスでセルフホストしたり、開発用にlocalhostでローカル実行したりできます。すぐに使える環境のカタログはopenreward.aiで利用できます。

このガイドでは、OpenRewardをTRLに統合する方法について説明します。標準自体の詳細については、ORSのドキュメントを参照してください。

統合はtrl.experimental.openrewardにあり、trl[openreward]エクストラによってゲートされています(遅延インポートのため、非ユーザーにはコストがかかりません)。

OpenReward環境を使用するタイミング

GRPOTrainer は、environment_factory スロットを介した環境ベースのトレーニングをサポートしています。一般的な契約については、OpenEnv を参照してください。ORS 対応の環境に対してトレーニングを行う場合は、OpenReward を使用します。具体的には、OpenReward カタログ(例: Eigent/SETA、kanishk/EndlessTerminals、nebius/SWE-rebench-V2)、独自のインフラでセルフホストする環境、または開発中のローカルサーバーなどが対象です。

インストール

pip install trl[openreward]

これにより、openreward Python SDK がインストールされます。統合自体は openreward を遅延インポートするため、trl.experimental.openreward を使用しないユーザーには影響しません。

クイックスタート

OpenRewardSpec クラスは、これらの kwarg 名と 1:1 で対応するプロパティを公開することで、単一の ORS 環境を TRL トレーナーの 3 つのスロット — train_dataset、environment_factory、reward_funcs — に接続します。

from trl import GRPOConfig, GRPOTrainer
from trl.experimental.openreward import OpenRewardSpec

spec = OpenRewardSpec("Eigent/SETA", num_tasks=64)

trainer = GRPOTrainer(
    model="Qwen/Qwen3-4B",
    args=GRPOConfig(
        num_generations=2,
        max_steps=5,
        max_tool_calling_iterations=20,
        log_completions=True,
    ),
    train_dataset=spec.train_dataset,
    environment_factory=spec.environment_factory,
    reward_funcs=spec.reward_funcs,
)
trainer.train()

OpenRewardSpec は内部で、初回アクセス時に遅延して3つの処理を実行します。

  1. spec.train_dataset: 環境のタスクリストからdatasets.Datasetを導出します(SDKを介したHTTPラウンドトリップが1回)。少なくともprompt、task_index、および各タスクのメタデータカラムが折り畳み込まれています。
  2. spec.environment_factory: 各呼び出しで新しいロールアウトごとのアダプターを生成する、引数なしの呼び出し可能オブジェクトを返します。アダプターは、ORS ツールごとに 1 つの Python メソッドを公開し、その型付きシグネチャとドキュメント文字列は、環境の JSON Schema から自動生成されます。TRL のツールコレクターは、inspect.getmembers を通じてそれらを収集します。
  3. spec.reward_funcs: 結果のみを対象とする報酬関数(軌跡内の最後の非 null 報酬)で、SETA などのスパース報酬環境に適しています。

ハブ環境の使用

ターゲットとしてopenreward.aiのカタログ名を渡します。SDKは認証のために環境からOPENREWARD_API_KEYを読み取ります。

spec = OpenRewardSpec("Eigent/SETA", num_tasks=64)

セルフホスト環境の使用

URLを直接渡します。サーバーがAPIキーを必須としていない場合、APIキーは不要です。

spec = OpenRewardSpec("https://my-org-my-env.hf.space", env_name="my_env")

openreward SDKは、デフォルトで2つのサブドメインを持つプラットフォームのレイアウトを想定しています(ステートレスな呼び出しにはapi.<host>、SSEベースのセッション呼び出しにはsessions.<host>)。単一ホストのセルフホスト型サーバー(すべてのリクエストを1つのURLで処理する構成)の場合、OpenRewardSpecを構築する前に、以下のオーバーライド用環境変数を設定してください。

import os

URL = "https://my-org-my-env.hf.space"
os.environ["OPENREWARD_API_URL"]     = URL
os.environ["OPENREWARD_SESSION_URL"] = URL

spec = OpenRewardSpec(URL, env_name="my_env")

ローカルで最小環境を実行する

外部依存なしで統合をエンドツーエンドで試す最も速い方法は、openreward SDKのEnvironment + Serverのスケルトンで定義される小さなORSサーバーです。以下の例は完全なecho環境です。モデルは、タスクのターゲット文字列を引数としてecho(text=...)を呼び出すことで勝利します。

# server.py
from pydantic import BaseModel
from openreward.environments import Environment, JSONObject, Server, TextBlock, ToolOutput, tool


class EchoTaskSpec(BaseModel):
    target: str

class EchoParams(BaseModel):
    text: str


class EchoEnvironment(Environment):
    def __init__(self, task_spec: JSONObject = {}, secrets: dict[str, str] = {}):
        super().__init__(task_spec)
        self.config = EchoTaskSpec.model_validate(task_spec)

    @classmethod
    def list_splits(cls) -> list[str]:
        return ["train"]

    @classmethod
    def list_tasks(cls, split: str) -> list[JSONObject]:
        return [{"target": "hello"}, {"target": "world"}]

    def get_prompt(self) -> list[TextBlock]:
        return [TextBlock(type="text", text=f"Echo '{self.config.target}' to win.")]

    @tool
    async def echo(self, params: EchoParams) -> ToolOutput:
        """Submit a string. Reward 1.0 + finished if it matches the target.

        Args:
            text: The string to echo back.
        """
        correct = params.text == self.config.target
        return ToolOutput(
            blocks=[TextBlock(type="text", text="match" if correct else "no match")],
            reward=1.0 if correct else 0.0,
            finished=correct,
        )


if __name__ == "__main__":
    Server([EchoEnvironment]).run(host="0.0.0.0", port=8000)

実行します:

pip install openreward fastapi uvicorn pydantic
python server.py     # listens on :8000

次に、OpenRewardSpecをそれに向けます(上記で説明したURLのオーバーライドを指定して):

import os
URL = "http://127.0.0.1:8000"
os.environ["OPENREWARD_API_URL"]     = URL
os.environ["OPENREWARD_SESSION_URL"] = URL

from trl.experimental.openreward import OpenRewardSpec
spec = OpenRewardSpec(URL, env_name="echoenvironment")
print(spec.train_dataset)        # 2 rows, task_index + target columns

これはTRLのテストでも使用されているフィクスチャパターンです。デプロイされたSpaceについては、trl-internal-testing/openreward-echo-envを参照してください。

タスクの選択

OpenRewardSpecは、カウントまたは明示的なインデックスリストのいずれかを指定できます。

spec = OpenRewardSpec("Eigent/SETA", num_tasks=10)                      # first 10 tasks
spec = OpenRewardSpec("Eigent/SETA", indices=[0, 5, 13, 27])            # specific indices
spec = OpenRewardSpec("Eigent/SETA", indices=list(range(50, 100)))      # range

num_tasksとindicesは排他的であり、どちらも必要なタスクのみを取得します(タスク全体のリストをスキャンしません)。

ツールバインディングの仕組み

構築時に、spec は env の /tools エンドポイントを呼び出し、ツール仕様のリスト(それぞれに名前、説明、引数の JSON Schema を持つ)を取得します。各ツールについて、ロールアウトごとのアダプターに、型付きのシグネチャとスキーマから導出された docstring を持つ Python メソッドを生成します。これらのメソッドはアダプターのクラスに設定され、TRL のツールコレクターによって検出されるため、transformers.utils.get_json_schema と TRL の inspect.getmembers(type(env), isfunction) の両方が、env ごとにラッパーコードを書くことなく、モデルに対して正しいツールスキーマを生成します。

ツール説明にPythonソースに安全に埋め込めない文字が含まれている場合、バインダーはサニタイズされた形式にフォールバックし、実際の環境でのバインディングが失敗しないようにします。

報酬関数

spec.reward_funcsのデフォルトは、アウトカムベースの報酬です。各ロールアウトに対して、軌跡中に観測された最後の非nullの報酬を返します。これは、報酬がスパースな環境(例えば、submit_solutionのみが非nullの報酬を返すSETAなど)に適したデフォルトです。

どのツールも null 以外の報酬を返さない場合、デフォルトの報酬関数は 0.0 ではなく None を返します。これには、ツール呼び出しがないロールアウト、スコアリング対象外のツール呼び出しのみを含むロールアウト、または失敗したツール呼び出しのみを含むロールアウトが含まれます。観測された 0.0 または負の報酬は、引き続き有効なスコアとして扱われます。後続のスコアリング対象外または失敗したツール呼び出しは、それ以前のスコアを消去しません。

GRPOでは、すべての報酬関数がNoneを返すロールアウトは報酬ベースラインから除外され、そのアドバンテージはゼロになります。これにより、損失正規化からそのトークンを除外したり、KL正則化などの他の項を無効にしたりすることはしません。他の報酬関数がそのロールアウトにスコアを付与する場合は、ゼロ以外のアドバンテージを受け取ることが依然として可能です。

スコアが欠落している場合、エピソードが終了した理由を特定できません。諦めやツール呼び出しの上限到達をタスクの失敗として扱うべき場合は、環境またはカスタム報酬関数から明示的な失敗報酬を返してください。インフラストラクチャやグレーダーの失敗は、スコア未算定として残すこともできます。平均報酬と並行して、スコアリングのカバレッジを監視してください。

カスタムな報酬を使用したい場合は、通常のTRL報酬関数を作成し、直接渡します:

def my_reward(environments, **kwargs) -> list[float | None]:
    rewards = spec.reward_funcs(environments=environments)
    return [2.0 * reward if reward is not None else None for reward in rewards]

trainer = GRPOTrainer(
    ...,
    reward_funcs=my_reward,
)

ロールアウトごとのアダプターは、env.reward、env.rewards、env.metadata、env.finished、env.last_output といった、TRL が必要とする実行状態を公開し、任意の事後報酬整形を可能にします。

OpenRewardSpec

class trl.experimental.openreward.OpenRewardSpec

< >

( target: strnum_tasks: int | None = Nonesplit: str = 'train'indices: list[int] | None = Noneapi_key: str | None = Nonesecrets: dict[str, str] | None = Noneenv_name: str | None = Noneinclude_metadata: bool = Truediscover_task_tools: bool = Truetask_tools_discovery_index: int | None = None )

パラメータ

  • target (str) — openreward.ai のカタログ名(“Eigent/SETA”)または、任意の ORS サーバーを指す URL(”https://you-seta.hf.space”、“http://localhost:8080”)。文字列内に :// が存在するかどうかで自動検出されます。
  • num_tasks (int、任意) — データセットに取り込むタスク数の上限。Noneの場合、envが公開するすべてのタスクを使用します。
  • split (str、optional、デフォルトは “train”) — タスクリストを取得するスプリットを指定します。
  • indices (list[int]、任意) — 学習に使用する特定のタスクのインデックス。num_tasks と排他です。デバッグやカリキュラムのサブセットに便利です。
  • api_key (str、任意) — OPENREWARD_API_KEY の上書き。 target がカタログ名の場合にのみ使用されます。
  • secrets (dict[str, str]、任意) — env.session(secrets=) に転送されるセッションごとのシークレット。
  • env_name (str, optional) — サーバー上で参照する環境名の上書きです。必要になることはほとんどありません。
  • include_metadata (bool、任意、デフォルトは True) — タスクごとのメタデータ(difficulty、category、tags など)をデータセットの行に折り込み、TRL の inputs 引数を通じて報酬関数がそれを読み取れるようにします。
  • discover_task_tools (bool、任意、デフォルトは True) — True の場合、短命な ORS セッションを開き、session.list_tools() を使用して、タスク固有のツール(ORS における GET …/task_tools、例えば @tool(shared=False) や list_task_tools())を GRPO にバインドします。 プローブが失敗した場合、environment.list_tools() のみにフォールバックします。 False に設定すると、追加のセッションをスキップします(共有ツールのみ / オフライン時の特殊な挙動)。
  • task_tools_discovery_index (int, optional) — 設定されている場合、ディスカバリーセッションでのみ使用されるタスクインデックスです。上書きして、以下のマルチインデックスプロービングを無効にします。省略され、indices= が設定されている場合、ディスカバリーはindices 内の各個別エントリ(ソート済み)ごとに1つのプロービングセッションを開き、マージしてツール仕様を名前ごとに統合し、列挙されたすべてのタスクのタスク固有のツールをバインドします。省略され、num_tasks / フルリストモードが使用されている場合、タスク0のみをプローブします。discover_task_tools=False の場合は無視されます。

TRL トレーナーに ORS 環境を接続するための単一の仕様オブジェクトです。

制限事項

  • 統合は trl.experimental にあります — API は変更される可能性があります。CI ログで警告を無効にするには、TRL_EXPERIMENTAL_SILENCE=1 を設定してください。
  • 現在、1つの環境をカバーする単一のOpenRewardSpecのみを公開しています。複数環境での学習(OpenEnvの「メタ環境」パターンに準拠)はまだサポートされていません。
  • 長時間のロールアウト(1エピソードあたり15分超)には、キープアライブのピンが必要です。まだ実装されていません。

リファレンス

Update on GitHub