TRL documentation

환경을 활용한 LLM 학습을 위한 OpenEnv 통합

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.18.0).
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 또는 컨테이너화 실행을 통해 환경을 백엔드 서버로 실행하는 것을 지원합니다. 바로 사용할 수 있는 OpenEnv 환경 모음은 Hugging Face Hub에서 찾을 수 있습니다.

이 가이드는 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"

이 명령은 원격 환경 서버와 WebSocket을 통해 통신하는 environment client(예: EchoEnv)와 action/observation 모델, 그리고 필요한 모든 의존성( openenv 포함)을 설치합니다.

모든 환경의 설치 명령은 해당 HF Space 페이지에서 찾을 수 있습니다. ⋮ (세 점) 메뉴를 클릭하고 “Use this Space”를 선택하면 설치 지침을 볼 수 있습니다.

pip install "openenv>=0.3.1"을 사용하여 PyPI에서 core 패키지도 설치할 수 있지만, 환경별 종속성은 별도로 설치해야 할 수 있습니다.

개발을 위해 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 environment를 사용하여 모델을 학습시키며, 이 환경은 텍스트 길이에 따라 보상을 부여합니다:

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())를 찾아내고 이를 함수 호출 도구로 노출합니다. 각 메서드는 타입이 지정된 인자를 포함한 적절한 docstring을 가져야 하며, 트레이너는 이를 기반으로 도구 스키마를 구성합니다.
  4. 멀티턴 루프: 트레이너가 completion을 생성하고, tool 호출을 파싱한 후 echo()를 실행하고, 결과를 추가한 뒤 다시 생성하는 과정을 반복합니다. 이 과정은 모델이 더 이상 tool을 호출하지 않거나 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 인수를 통해 인터랙티브 환경 학습을 지원합니다. 이 인수가 제공되면 트레이너는 멀티턴 도구 호출 루프를 자동으로 처리합니다. 즉, 완성본을 생성하고, 도구 호출을 파싱하며, 환경에서 이를 실행한 후 결과를 모델에 다시 전달합니다. 커스텀 롤아웃 코드 없이도 가능합니다.

Environment 클래스 요구 사항

환경 클래스는 다음 규칙을 따라야 합니다:

  • __init__(self) (선택 사항): 제공되는 경우, 인자를 받지 않아야 합니다. 상태 또는 클라이언트를 초기화하는 데 사용하세요. 외부 구성(예: URL)이 필요한 경우, 외부 범위 또는 모듈 수준 변수에서 이를 가져와야 합니다.
  • reset(self, **kwargs): 각 에피소드 시작 시 호출됩니다. 모든 데이터셋 컬럼을 키워드 인수로 받습니다. 문자열 형태의 관측값을 반환합니다(초기 관측값이 없으면 None을 반환).
  • Tool methods: reset를 제외한 모든 public 메서드(즉, _로 시작하지 않는 메서드)는 자동으로 tool로 노출됩니다. trainer가 모델용 tool schema를 생성할 때 이를 사용하기 때문에, 각 tool 메서드는 Args: 설명이 포함된 docstring을 가져야 합니다.

환경 클래스 팁

  • 보상 상태: 환경 인스턴스에 원하는 상태(예: self.reward, self.done 등)를 저장하고, environments 매개변수를 통해 보상 함수에서 이를 접근할 수 있습니다. 이 패턴의 예시는 빠른 시작 가이드를 참고하세요.
  • 오류 처리: 도구 메서드가 예외를 발생시키면(예: ValueError("Game over.")), trainer가 이를 포착하여 오류 메시지를 도구 응답으로 모델에 전달합니다. 이는 동작이 유효하지 않거나 에피소드가 종료되었음을 알리는 권장 방식입니다.
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"

Tools는 개별 메서드여야 하며, 설명적인 이름과 타입이 지정된 인자(예: guess(word: str), move(direction: str))를 가져야 합니다. 모델이 tool calling을 학습하려면 의미 있는 tool 이름과 인자 설명이 필요하므로, step(action)와 같은 범용 메서드를 사용하는 것은 권장하지 않습니다.

보상 함수

environments 매개변수(환경 인스턴스 목록)를 받는 보상 함수는 에피소드 동안 저장된 모든 상태에 접근할 수 있습니다:

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

보상 함수에 대한 자세한 내용은 GRPO - Custom Reward Functions를 참조하세요.

보상 함수를 위한 팁

OpenEnv 환경과 GRPO를 함께 작업할 때 유용했던 몇 가지 사항을 공유합니다:

  • 그룹이 서로 다름을 확인하세요. GRPO는 그룹 내에서 완성 결과를 비교하므로 절대값보다 상대 순위가 더 중요하며, 모든 완성 결과가 동일한 보상을 받는 그룹은 어떤 신호도 제공하지 않습니다. 모델이 때로는 성공하는 경우, 이진 보상(성공 시 1.0, 그 외 0.0)이 깔끔합니다. 성공률이 낮을 때는 부분 점수를 추가하세요. 스도쿠 예제는 풀린 보상에 채워진 칸의 비율을 더하고, 워들 예제는 가장 좋은 추측의 그린 수를 더합니다.
  • 경로가 아닌 최종 상태를 확인하세요. 가능하면 모델이 특정 동작 순서를 따랐는지 확인하는 대신, 환경이 결과를 판단하도록 하세요(예: “모델이 퍼즐을 풀었는가?”). 이렇게 하면 모델이 자신만의 전략을 발견할 자유가 생깁니다.
  • 학습 전에 보상 함수를 테스트하세요. 수동으로 몇 에피소드를 실행하여(예: Wordle 예제 노트북) 환경이 적절한 보상을 반환하는지 확인하세요. 성능이 좋은 모델조차 무작위 베이스라인보다 높은 점수를 얻지 못한다면, 보상 신호를 조정해야 할 수 있습니다.

멀티턴 에피소드에서 max_completion_length

max_completion_length 매개변수는 단일 생성이 아니라 멀티턴 대화 전체에 걸친 토큰 총 개수(모든 모델 생성 + 도구 결과 합산)를 제한합니다. 턴 수가 많은 환경(예: 수십 개의 수를 포함하는 Sudoku)에서는 이 값을 늘려야 할 수 있습니다:

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

에피소드가 중간에 끊기는 경우(모델이 게임 도중 중단되는 경우)에는 이것이 원인이 될 가능성이 높습니다.

Space에 대한 긴 실행

Space는 몇 시간 동안 열려 있는 WebSockets를 닫을 수 있습니다(약 2시간마다 발생하는 것을 확인했습니다). 0.11.0 이하의 OpenEnv 클라이언트는 새로운 환경 세션에 다시 연결하며, 게임이 진행 중이었다면 다음 step()은 클라이언트가 리셋하지 않은 게임에서 오류 없이 실행됩니다. OpenEnv 0.12.0부터 클라이언트는 대신 SessionLostError를 발생시켜 게임이 실패합니다. 학습 과정에서 이러한 연결 끊김을 방지하려면 reset()에서 새 클라이언트를 열어 각 연결이 에피소드 하나만 지속되도록 하고, 이전 세션이 비동기로 해제되므로 동시에 실행되는 게임 수보다 약간 큰 값을 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 단계 및 평가, 그리고 하드웨어와 비용을 포함한 전체 파이프라인을 다룹니다. 이 예제의 노트북 버전은 grpo_wordle.ipynb입니다.

TextArena 환경

TextArena는 Wordle, Snake, Tic-Tac-Toe와 같은 텍스트 기반 게임을 활용하여 LLM의 추론 능력을 평가하도록 설계된 오픈 소스 경쟁 텍스트 기반 게임 모음입니다.

image of TextArena

왜 Wordle인가요?

Wordle은 피드백에 대한 추론이 필요하고, 순수하게 텍스트 기반이며, 1B 파라미터 규모의 모델도 이를 통해 성능을 향상시킬 수 있기 때문에 환경 기반 RL의 좋은 벤치마크입니다.

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개의 게임에서 1에포크 동안 학습된 모델입니다. 각 추측 전에 짧은 추론이 포함됩니다.

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는 매 턴마다 전체 게임 히스토리를 반환합니다. 컨텍스트를 반복하지 않도록 새로운 부분만 슬라이싱하여 추출합니다.
  • 완료 시 예외: 게임이 종료된 후에도 모델이 추측을 시도하면 guess()가 ValueError를 발생시킵니다. 트레이너는 이를 포착하여 "Game over."를 도구 응답으로 모델에 다시 전달하고, tools/failure_frequency에 이를 집계합니다.

보상 함수 및 학습

환경 자체의 보상(reward)은 1번의 추측으로 이기든 6번으로 이기든 동일한 승리를 지급하며, 부분 점수는 게임이 끝날 때만 지급합니다. 사고(thinking)가 켜져 있는 경우, 정책(policy)은 사고를 줄이고 6번의 추측을 모두 사용해 게임의 끝까지 도달하는 것을 학습할 수 있습니다. 이를 방지하기 위해 세 가지 보상 항(reward term)을 사용합니다:

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회 추측 시 1.42, 6회 시 1.0). 패배는 게임이 종료되었든 아니든 게임 내 최선의 추측에 대한 보상을 지급하며, 각 잘못된 이동은 0.1의 비용을 발생시킵니다. TRL은 각 항을 별도로 기록하므로 rewards/reward_win/mean은 승률과 함께 증가합니다.

예제 실행

각 단계에서 64개의 게임을 플레이하며, 게임당 환경 세션이 하나씩 생성됩니다. openenv/wordle Space를 복제하고, 복제본에서 MAX_CONCURRENT_ENVS=64 변수를 설정하세요.

colocate
서버

Colocate 모드 (GPU 1개, 권장)

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

이 방법은 학습과 동일한 프로세스에서 vLLM을 실행하므로 단일 GPU만 필요합니다. 저희 실행에서는 A100 80GB 1대에서 150단계에 약 2시간 15분이 소요되었습니다.

결과

…

모델승승리당 추측 횟수잘못된 추측
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은 None 값을 nansum/nanmean으로 처리합니다.

예: 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. ..."""
        ...

주요 패턴:

  • 지연된 클라이언트 초기화: 불필요한 WebSocket 연결을 피하기 위해 클라이언트를 __init__()가 아닌 reset()에서 생성합니다.
  • 다시 열기 전에 닫기: 서버 용량 오류를 방지하려면 이전 클라이언트를 닫은 후 새 클라이언트를 생성하세요.
  • 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을 집계 연산에 사용하므로, 각 샘플은 해당되는 보상 함수로만 점수를 받습니다.

환경 라우팅이 포함된 Dataset

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가 아니라 보상 함수별 지표(train/reward_func_0, train/reward_func_1 등)를 모니터링하세요. 합산 지표는 환경 사이를 오가며 변동하기 때문에 노이즈가 있는 것처럼 보일 수 있습니다.

환경 실행

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 연결(생성당 1개)을 열며, 공유 Space는 이를 지원하지 않을 수 있습니다. 자세한 내용은 Server concurrency를 참고하세요.

환경 카탈로그

유지 관리 중인 환경의 현재 카탈로그를 탐색하는 가장 좋은 방법은 공식 OpenEnv 카탈로그를 방문하는 것입니다.

자체 환경을 만들려면 OpenEnv로 자체 환경 구축하기 가이드를 참고하세요. 환경은 Hub와 긴밀하게 통합되어 있어, 커뮤니티가 재사용할 수 있도록 새로운 환경을 push할 수 있습니다.

서버 동시성

environment_factory를 사용할 때, trainer는 N개의 환경 인스턴스(생성당 1개)를 생성하며, 각 인스턴스는 서버에 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이고 배치 크기가 1인 경우, 최소 64개의 동시 세션이 필요합니다.

environment_factory와 rollout_func

GRPOTrainer는 환경 기반 학습을 위한 두 가지 접근 방식을 지원합니다:

  • environment_factory (권장): 도구 메서드를 포함한 환경 클래스를 정의하면, 트레이너가 생성, 도구 호출 파싱, 다중 턴 루프를 자동으로 처리합니다. 이 가이드 전체에서 이 방식을 사용합니다.
  • rollout_func: 생성 및 환경 상호작용 루프를 직접 작성합니다. 이를 통해 완성본 생성 방식, 도구 실행 방식, 보상 계산 방식을 완전히 제어할 수 있습니다.

environment_factory가 사용 사례에 맞지 않는 경우 rollout_func를 사용하세요. 예를 들어, 외부 서버가 생성 루프를 소유하고 자체 에이전트-환경 상호작용 프로토콜을 관리하는 외부 에이전트 서버가 있습니다.

하네스에서의 학습: 실제 코딩 에이전트 학습(Harbor)

위 통합들은 white-box 방식입니다. TRL이 멀티턴 루프를 직접 제어합니다. 각 턴을 샘플링하고, 도구 호출을 파싱한 후 실행하며, 결과를 다시 전달합니다.

일부 에이전트는 자체 루프를 소유하고 있어 이러한 방식으로 구동할 수 없습니다. opencode나 claude-code와 같은 프로덕션 코딩 에이전트 하네스는 자체 플래너, 도구 세트, 컨텍스트 관리, 종료 조건을 갖추고 있습니다. 해당 에이전트의 재구현이 아니라, 그 에이전트 자체를 학습시키고 싶을 것입니다.

이를 위해 TRL은 AsyncGRPOTrainer를 기반으로 구축된 실험적인 블랙박스(루프 소유) 경로를 제공합니다. HarnessAsyncGRPOTrainer는 OpenEnv 세션 팩토리를 구동합니다. 에이전트와 해당 샌드박스는 Harbor에서 제공되며, openenv harbor serve를 통해 OpenEnv로 서빙됩니다. 따라서 Harbor가 설치하는 모든 에이전트는 동일한 방식으로 학습할 수 있으며, --harness를 사용하여 하나를 선택합니다. Hugging Face Jobs 런처를 포함한 완전한 예제는 examples/async_grpo_harbor를 참고하세요.

이 경로는 TrainingTrace 캡처 API가 포함된 첫 릴리스인 openenv>=0.7.0이 필요합니다: pip install "openenv>=0.7.0".

작동 방식

TRL은 여기서 각 턴을 샘플링하지 않습니다. 에이전트가 스스로 완료까지 실행되고, TRL은 에이전트가 수행한 내용을 다시 읽어옵니다:

  1. OpenEnv Harbor 서버는 샌드박스(예: E2B)에서 Harbor 작업을 수행하는 에이전트를 시작합니다. 에이전트가 수행하는 모든 모델 호출은 서버의 캡처 프록시를 통해 vLLM 서버로 전달되며, 이 서버는 엔진 자체의 프롬프트 및 완성 토큰 ID, 샘플링된 logprobs, 토큰별 손실 마스크를 기록합니다.
  2. 에이전트가 중단되면, 해당 작업 자체의 검증기가 결과를 평가하고 TRL은 OpenEnv에서 검증된 TrainingTrace를 받아 그 토큰 ID와 마스크로부터 학습 행을 구성합니다. TRL은 캡처된 프롬프트를 다시 토큰화하지 않습니다.
  3. GRPO는 해당 행을 기준으로 학습합니다. 보상은 그룹 상대적 이점을 통해 학습된 모든 토큰으로 전파되고, 새로운 가중치는 동일한 vLLM 서버로 동기화됩니다.

각 rollout는 자체 격리된 세션에서 실행되며, 전용 샌드박스와 캡처 세션을 가집니다. 이러한 격리는 두 가지 이유로 중요합니다. 프록시는 해당 rollout의 토큰을 정확히 캡처해야 하고, 한 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를 반환합니다. 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은 캡처 검증 및 토큰 선택을 담당합니다. TRL은 프로듀서의 전체 시퀀스 마스크를 내부 completion 마스크로 변환하고, 결과된 행을 패킹한 뒤 GRPO를 적용합니다. 토큰을 재구성하거나 어떤 호출이 에이전트에 속하는지 추론하지는 않습니다.

rollout_reward_fn : rollout를 어떻게 점수화할지

Callable[[HarnessRolloutOutcome], float | None]

TRL은 사용자의 작업에서 성공이 무엇을 의미하는지 알 수 없으므로, 완료된 각 rollout을 스칼라 리워드로 변환합니다. 이 함수는 에이전트가 수행한 작업을 설명하는 HarnessRolloutOutcome를 인수로 받습니다:

  • env_reward (float | None): 작업 검증기(verifier)에서 반환된 보상(reward)이며, 롤아웃(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을 반환하여 롤아웃을 점수화할 수 없는 것으로 표시합니다. 이렇게 하면 해당 롤아웃은 0으로 계산되는 대신 그룹 베이스라인에서 제외됩니다. 이 함수를 전달하지 않으면, 원본 env_reward가 그대로 사용됩니다. Harbor 예제는 검증기의 정확도를 주요 신호로 유지하고, 작업이 해결된 경우에만 작은 도구 효율 보너스를 추가합니다.

왜 검증기만 사용하지 않는가? 검증기는 환경의 역할이며, “결과가 얼마나 정확했는가”라는 하나의 질문에 답함으로써 평가에 깨끗하고 재사용 가능하게 유지됩니다. 학습에 사용하는 보상(리워드)는 점수를 이진화하거나, 퇴화한 행동을 벌점 처리하거나, 점수를 매길 수 없는 롤아웃을 제외하는 등 학습 시점의 독립적인 결정입니다. 또한 검증기는 최종 결과만 검사하기 때문에, rollout_reward_fn이 받는 신호(트라우제리, 도구 호출 횟수, timed_out, 트레이스)는 검증기가 절대 볼 수 없습니다. 예를 들어, 어떤 롤아웃은 일부 테스트를 통과하되 bash를 한 번도 실행하지 않을 수 있습니다. 이를 보고 벌점을 부과할 수 있는 것은 rollout_reward_fn뿐입니다.

Token eligibility는 OpenEnv에서 가져옵니다

session.fetch_training_trace()는 검증된 TrainingTrace를 반환합니다. 각 턴에는 고유한 node_id, prompt_token_ids, completion_token_ids, per_token_logps가 있으며, 프롬프트와 completion을 아우르는 필수 loss_mask가 있습니다. 프롬프트 위치의 값은 0입니다. 어시스턴트 응답 내부의 0을 포함한 부분 completion 마스크는 그대로 유지됩니다.

예를 들어, 프롬프트 [10, 11], 완성 [20, 21, 22] 및 마스크 [0, 0, 1, 0, 1]는 토큰 20과 22에 대해 학습합니다. 토큰 21은 컨텍스트로 유지됩니다. 완전히 마스크된 완성은 컨텍스트 및 사용량 계산을 위해 유지됩니다.

producer는 보조 호출과 폐기된 재시도를 제외합니다. 이 API에는 train_turn_fn 또는 agent_turn_fn 매개변수가 없습니다. 기존 통합은 선택 로직을 OpenEnv 마스크로 이동해야 합니다. 작업의 검증기는 토큰 자격과 독립적으로 유지됩니다.

잘못된 포맷의 캡처는 OpenEnv 경계에서 실패하여 학습을 중단합니다. 전송 실패는 여전히 점수를 매길 수 없습니다. 다시 작성된 히스토리는 별도의 행을 시작하며, 샘플링된 토큰과 logprobs를 보존합니다. 여러 행은 여전히 rollout reward를 공유하며, 이 API는 GRPO의 가중치를 변경하지 않습니다.

Metrics

harness trainer는 AsyncGRPOTrainer의 tool metrics를 동일한 이름으로 기록하며, 다음과 같은 차이가 있습니다:

  • 에이전트가 루프를 소유하는 경우(harness_adapter 없음)에는 마지막 요청의 도구 결과를 기반으로 실패를 추정합니다. 결과의 텍스트에 error, failed, traceback 또는 exception가 포함된 경우(대소문자 구분 없음) 실패로 간주하며, 동일한 도구와 동일한 첫 200자리를 가진 결과는 한 번만 집계됩니다. OpenEnv의 트레이스에 타임스탬프가 없으므로 tools/latency_s와 tools/<name>/latency_s은 기록되지 않습니다.
  • harness_adapter를 사용하는 경우, harness가 해당 도구 호출에 대해 오류를 보고하면 실패로 간주됩니다. 한 번의 모델 응답과 다음 요청 사이의 시간은 해당 턴의 도구 호출에 걸쳐 균등하게 분배되어 tools/latency_s로 기록되며, 이는 harness 자체의 오버헤드를 포함합니다. 마지막 턴은 측정되지 않습니다.
  • tools/unknown_name_count, tools/parallel_calls_mean, tools/loop_exhausted_frac은 내장된 tool loop에서 생성되며 기록되지 않습니다.
  • 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 서버는 tool calling과 실제 token id를 노출하고, NCCL weight sync를 활성화하여 에이전트가 항상 최신 policy를 사용하도록 해야 합니다. Rollouts와 weight 업데이트는 동일한 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 training은 trl.experimental 아래에 있으며, 해당 API는 변경될 수 있습니다. 예제는 테스트에 사용된 OpenEnv 리비전을 고정하고 있으며, 정확한 설정은 README를 참고하세요.

Update on GitHub