TRL documentation

Integração do OpenEnv para Treinar LLMs com Ambientes

TRL

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).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Integração do OpenEnv para Treinar LLMs com Ambientes

OpenEnv é um framework de código aberto para definir, implantar e interagir com ambientes em aprendizado por reforço (RL) e fluxos de trabalho agênticos. Ele fornece APIs padronizadas para interação com ambientes e suporta a execução de ambientes como servidores de backend (via WebSocket ou execução em contêineres). Você pode encontrar uma coleção de ambientes OpenEnv prontos para uso no Hugging Face Hub.

Este guia aborda como integrar o OpenEnv com o TRL. Para mais informações sobre o próprio OpenEnv, consulte a documentação do OpenEnv.

Exemplos prontos de OpenEnv: grpo_echo (mínimo), grpo_catch, grpo_wordle, grpo_sudoku, grpo_multi_env, grpo_browsergym, grpo_carla e async_grpo_harbor (local e HF Jobs).

Quando usar ambientes

GRPOTrainer pode ser usado para treinar agentes. Para tarefas agênticas, ele suporta dois modos: tools, em que o modelo pode chamar funções externas, mas cada chamada é sem estado e independente, e environments, que mantêm estado entre as rodadas, permitindo interação genuína de múltiplas rodadas em que as ações do agente moldam as observações futuras. Use environments quando a continuidade for importante — por exemplo, ao navegar em um jogo, ao visualizar uma página da web ou em qualquer tarefa em que o que o agente vê a seguir depende do que ele fez antes.

Escolhendo uma integração de ambiente

OpenEnv é o caminho nativo documentado aqui. Duas integrações adicionais — OpenReward e Harbor — seguem o mesmo contrato environment_factory e são intercambiáveis no nível do TRL. Consulte a comparação das integrações de ambiente no guia do GRPO para escolher aquela cujo ecossistema se adequa à sua tarefa.

Instalação

Os ambientes OpenEnv são hospedados como Spaces do Hugging Face, que também são repositórios Git instaláveis via pip:

# 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"

Isso instala o cliente de ambiente (por exemplo, EchoEnv) que se comunica com o servidor de ambiente remoto via WebSocket, juntamente com os modelos de ação/observação e todas as dependências necessárias (incluindo openenv).

Você pode encontrar o comando de instalação para qualquer ambiente na página do seu HF Space. Clique no menu ⋮ (três pontos) e selecione “Use this Space” para ver as instruções de instalação.

Você também pode instalar o pacote principal do PyPI com pip install "openenv>=0.3.1", mas observe que as dependências específicas do ambiente podem precisar ser instaladas separadamente.

Para desenvolvimento, você pode clonar o repositório OpenEnv e instalar localmente:

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

Cada script de ambiente no TRL inclui metadados de dependências inline (PEP 723), permitindo que você também os execute diretamente com uv:
uv run examples/grpo_echo/grpo_echo.py
Isso instala automaticamente o pacote de ambiente necessário em um ambiente virtual isolado.

Início rápido

A maneira mais rápida de entender a integração é um exemplo completo. O script echo.py treina um modelo com o Echo environment, que recompensa as conclusões com base no comprimento do texto:

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

from trl import GRPOConfig, GRPOTrainer

ENV_URL = "https://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()

É isso. Veja o que acontece por baixo dos panos:

environment_factory=EchoToolEnv: O trainer cria uma instância de EchoToolEnv por geração (passe a classe, não uma instância).
reset() é chamado no início de cada episódio para inicializar o estado. Retorna uma string de observação (ou None).
Descoberta de ferramentas: O trainer descobre todos os métodos públicos na instância do ambiente (aqui, echo()) e os expõe como ferramentas de chamada de função. Cada método deve ter uma docstring adequada com argumentos tipados, que o trainer usa para construir o esquema da ferramenta.
Laço de múltiplas rodadas: o trainer gera uma conclusão, analisa as chamadas de ferramentas, executa echo(), adiciona o resultado e gera novamente, até que o modelo pare de chamar ferramentas ou max_completion_length seja atingido.
Função de recompensa: Lê env.reward de cada instância de ambiente após o episódio (antes de o ambiente ser reiniciado).

# 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://openenv-echo-env.hf.space

Abaixo está a curva de recompensa do treinamento:

Você pode explorar mais exemplos prontos de OpenEnv na nota no topo desta página.

Como o environment_factory funciona

O GRPOTrainer do TRL suporta treinamento em ambientes interativos por meio do argumento environment_factory. Quando fornecido, o trainer gerencia automaticamente o loop de chamadas de ferramentas em múltiplas rodadas: ele gera completions, analisa as chamadas de ferramentas, as executa no ambiente e devolve os resultados ao modelo. Tudo isso sem código de rollout personalizado.

Requisitos da classe Environment

Sua classe de ambiente deve seguir estas regras:

__init__(self) (opcional): Se fornecido, não deve aceitar argumentos. Use-o para inicializar o estado ou os clientes. Se você precisar de configuração externa (por exemplo, uma URL), capture-a do escopo circundante ou de variáveis de nível de módulo.
reset(self, **kwargs): Chamado no início de cada episódio. Recebe todas as colunas do dataset como argumentos de palavra-chave. Retorna uma observação em string (ou None para não haver observação inicial).
Métodos de ferramenta: qualquer método público (que não comece com _), exceto reset, é exposto automaticamente como uma ferramenta. Cada método de ferramenta deve ter uma docstring com descrições de Args:, pois o trainer usa essas informações para gerar o schema da ferramenta para o modelo.

Dicas para classes de ambiente

Estado para a recompensa: Você pode armazenar qualquer estado que desejar na instância do ambiente (por exemplo, self.reward, self.done, etc.) e acessá-lo na sua função de recompensa por meio do parâmetro environments. Consulte o guia de início rápido para um exemplo desse padrão.
Tratamento de erros: se um método de ferramenta lança uma exceção (por exemplo, ValueError("Game over.")), o trainer a captura e devolve a mensagem de erro ao modelo como resposta da ferramenta. Esta é a maneira recomendada de sinalizar que uma ação é inválida ou que o episódio terminou.

ENV_URL = "https://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"

As ferramentas devem ser métodos individuais com nomes descritivos e argumentos tipados (por exemplo, guess(word: str), move(direction: str)). Não recomendamos o uso de métodos genéricos como step(action), pois o modelo precisa de nomes de ferramentas significativos e descrições de argumentos para aprender a chamada de ferramentas.

Funções de recompensa

As funções de recompensa recebem o parâmetro environments (uma lista de instâncias de ambiente), então você pode acessar qualquer estado armazenado durante o episódio:

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

Para mais informações sobre funções de recompensa, consulte o GRPO - Funções de recompensa personalizadas.

Dicas para funções de recompensa

Algumas coisas que nos ajudaram ao trabalhar com ambientes OpenEnv e GRPO:

Verifique se os grupos diferem. O GRPO compara as conclusões dentro de um grupo, portanto, o ranking relativo é mais importante do que os valores absolutos, e um grupo em que todas as conclusões recebem a mesma recompensa não fornece nenhum sinal. Uma recompensa binária (1,0 para sucesso, 0,0 caso contrário) é limpa quando o modelo tem sucesso em alguns casos. Quando isso raramente ocorre, adicione crédito parcial: o exemplo de Sudoku adiciona a fração de células preenchidas à recompensa de resolvido, e o exemplo de Wordle adiciona os verdes da melhor tentativa.
Verifique o estado final, não o caminho. Sempre que possível, deixe que o ambiente avalie o resultado (por exemplo, “o modelo resolveu o quebra-cabeça?”) em vez de verificar se ele seguiu uma sequência específica de ações. Isso dá ao modelo a liberdade de descobrir suas próprias estratégias.
Teste sua recompensa antes do treinamento. Execute alguns episódios manualmente (veja o notebook de exemplo do Wordle) para confirmar que o ambiente retorna recompensas sensatas. Se um modelo capaz não consegue pontuar mais alto que uma linha de base aleatória, o sinal de recompensa pode precisar de ajuste.

max_completion_length em episódios de múltiplas rodadas

O parâmetro max_completion_length limita o número total de tokens em toda a conversa de múltiplas rodadas (todas as gerações do modelo + resultados das ferramentas combinados), e não apenas uma única geração. Para ambientes com muitas rodadas (por exemplo, Sudoku com dezenas de jogadas), você pode precisar aumentá-lo:

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

Se os episódios estão sendo encerrados prematuramente (o modelo para no meio do jogo), essa é provavelmente a causa.

Execuções longas contra um Space

Um Space pode fechar WebSockets que permanecem abertos por horas (observamos isso a cada 2 horas aproximadamente). Clientes do OpenEnv até a versão 0.11.0 então se reconectam a uma nova sessão de ambiente e, se um jogo estivesse em andamento, o próximo step() é executado em um jogo que o cliente nunca reiniciou, sem nenhum erro. Desde a versão 0.12.0 do OpenEnv, os clientes lançam SessionLostError em vez disso, fazendo com que o jogo falhe. Para evitar que essas quedas afetem o treinamento, abra um novo cliente em reset(), para que cada conexão dure um episódio, e defina MAX_CONCURRENT_ENVS um pouco acima do número de jogos executados simultaneamente, pois as sessões antigas são liberadas de forma assíncrona:

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

Exemplo avançado: Wordle

Vamos treinar um modelo para jogar Wordle usando o ambiente TextArena. Isso demonstra interação de múltiplas rodadas, processamento de feedback cumulativo, uma recompensa dividida em vários termos e término de episódio por meio de exceções.

O tutorial OpenEnv Wordle abrange todo o pipeline: os dados de SFT, o aquecimento de SFT, esta etapa de GRPO e a avaliação, com hardware e custos. A versão em notebook deste exemplo é grpo_wordle.ipynb.

O Ambiente TextArena

TextArena é uma coleção de jogos competitivos baseados em texto, de código aberto, projetada para avaliar as habilidades de raciocínio em LLMs usando jogos textuais como Wordle, Snake, Jogo da Velha e outros.

image of TextArena

Por que Wordle?

Wordle é um bom benchmark para RL baseado em ambiente, pois exige raciocínio sobre o feedback, é puramente baseado em texto e modelos com 1B de parâmetros podem melhorar no jogo.

Como o Wordle funciona? O Wordle é um jogo de adivinhação de palavras em que o jogador deve adivinhar uma palavra de 5 letras em 6 tentativas. Após cada tentativa, o ambiente fornece feedback letra por letra:
G U E S S
X G Y X X
X = não está na palavra, G = posição correta (verde), Y = posição incorreta (amarelo). Aqui, “U” está correto e no lugar certo, “E” está na palavra, mas em posição errada.

Comece com um aquecimento de SFT

O Qwen3-1.7B quase nunca vence uma partida, então o GRPO a partir do modelo base tem pouco para aprender: um grupo de partidas em que nenhuma é vencida não fornece nenhum sinal. O exemplo parte de openenv/Qwen3-1.7B-Wordle-SFT, treinado por uma época em 1.000 partidas jogadas por um solver no mesmo formato de tool-calling, com um raciocínio curto antes de cada palpite.

Classe Environment

A classe WordleEnv encapsula o cliente TextArena, expõe guess() como a ferramenta e mantém o que a recompensa precisa:

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

Principais escolhas de design:

reset() retorna a mensagem inicial do jogo como a primeira observação que o modelo vê.
guess() é a única ferramenta. O modelo a chama a cada turno com uma palavra de 5 letras.
Fatiamento de feedback cumulativo: o TextArena retorna o histórico completo do jogo a cada turno. Fatiamos apenas a parte nova para evitar repetir o contexto.
Exceção em done: Se o modelo tentar fazer uma suposição após o fim do jogo, guess() levanta um ValueError. O trainer captura isso e devolve "Game over." ao modelo como resposta da ferramenta, e o conta em tools/failure_frequency.

Funções de recompensa e treinamento

A própria recompensa do ambiente paga um valor igual por uma vitória, seja após 1 ou 6 tentativas, e concede crédito parcial apenas no final de um jogo. Com o raciocínio ativado, uma política pode aprender a pensar menos e usar todas as 6 tentativas para chegar ao final. Três termos de recompensa evitam isso:

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()

Uma vitória paga mais quanto mais rápido ela ocorre (1,42 em uma tentativa, 1,0 em seis), uma derrota paga pelo melhor palpite do jogo, independentemente de ele ter terminado ou não, e cada movimento inválido custa 0,1. O TRL registra cada termo separadamente, portanto rewards/reward_win/mean aumenta com a taxa de vitórias.

Executando o exemplo

Cada etapa joga 64 partidas, uma sessão de ambiente cada. Duplique o Space openenv/wordle e defina a variável MAX_CONCURRENT_ENVS=64 na sua cópia.

colocate

server

Resultados

Na nossa execução, com uma execução de treinamento por modelo, nas mesmas 30 palavras geradas de forma gulosa:

Model	Vitórias	Chutes por vitória	Chutes inválidos
Qwen3-1.7B	0/30	-	24
+ SFT	10/30	3.2	41
+ SFT + GRPO (modelo)	20/30	4.1	3

O GRPO eliminou a maioria das palavras repetidas e inventadas e venceu mais jogos. A maioria das novas vitórias levou de 4 a 6 tentativas, então o número de tentativas por vitória aumentou. As 30 palavras vêm dos 156 substantivos dos quais o TextArena extrai seus segredos, que os modelos também viram no treinamento, e uma nova execução resultará em algumas vitórias a mais (duas execuções de SFT com a mesma receita venceram 10 e 13 jogos). Veja os três modelos jogando na Wordle Arena e observe as curvas de treinamento no Trackio.

Treinamento em múltiplos ambientes

Você pode treinar um único modelo em vários ambientes simultaneamente. Isso é útil quando você deseja que um modelo aprenda diferentes habilidades em paralelo. Por exemplo, jogar Wordle (raciocínio linguístico) e Catch (raciocínio espacial) na mesma execução de treinamento.

A ideia central é criar uma classe de meta-ambiente que encapsula múltiplos ambientes e roteia cada amostra para o ambiente correto usando uma coluna do conjunto de dados.

Como funciona

Adicione uma coluna "env" (ou similar) ao seu conjunto de dados que identifique a qual ambiente cada amoça pertence.
Em reset(**kwargs), leia kwargs["env"] para selecionar o ambiente ativo para aquele episódio.
Exponha as ferramentas de todos os ambientes; o trainer descobre todos os métodos públicos.
Use funções de recompensa separadas por ambiente, retornando None para amostras que não pertencem a esse ambiente. O TRL trata valores None com nansum/nanmean.

Exemplo: Wordle + Catch

O script grpo_multi_env.py treina simultaneamente em Wordle e 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. ..."""
        ...

Principais padrões:

Inicialização preguiçosa do cliente: Crie os clientes em reset(), não em __init__(), para evitar conexões WebSocket desnecessárias.
Fechar antes de reabrir: Feche o cliente anterior antes de criar um novo para evitar erros de capacidade do servidor.
roteamento de kwargs: a coluna "env" do conjunto de dados é passada para reset() como um argumento de palavra-chave.
Todas as ferramentas são expostas simultaneamente: o modelo vê guess, move e stay como ferramentas disponíveis, independentemente do ambiente ativo. Se ele chamar a ferramenta errada (por exemplo, move durante o Wordle), o método levanta um ValueError que o trainer captura de forma graciosa. Na prática, os modelos aprendem a usar as ferramentas corretas com base no system prompt.

Funções de recompensa por ambiente

Cada função de recompensa retorna None para amostras de outros ambientes:

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

O TRL converte None para nan internamente e usa nansum/nanmean para agregação, de modo que cada amostra é pontuada apenas pela sua função de recompensa relevante.

Conjunto de dados com roteamento de ambiente

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

Executando o exemplo de multi-ambiente

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

Ao treinar em vários ambientes, monitore as métricas por função de recompensa (train/reward_func_0, train/reward_func_1, etc.) em vez da métrica combinada train/reward. A métrica combinada alterna entre os ambientes e pode parecer ruidosa.

Executando os ambientes

Ao usar environment_factory, o trainer se conecta ao servidor de ambiente automaticamente. Você só precisa que o servidor esteja em execução. Existem três maneiras de executar um servidor de ambiente OpenEnv:

space

docker

local

Catálogo de ambientes

A melhor maneira de explorar o catálogo atual de ambientes mantidos é visitar o catálogo oficial do OpenEnv.

Para criar seu próprio ambiente, consulte o guia sobre Building Your Own Environment with OpenEnv. Os ambientes estão fortemente integrados ao Hub, permitindo que você envie novos ambientes para que a comunidade os reutilize.

Concorrência do servidor

Ao usar environment_factory, o trainer cria N instâncias de ambiente (uma por geração), cada uma abrindo uma conexão WebSocket com o servidor. Por padrão, os servidores OpenEnv permitem apenas 1 sessão concorrente, o que causará falhas durante o treinamento.

Para suportar o treinamento paralelo, configure o servidor para concorrência:

No arquivo de ambiente, declare o suporte a sessões concorrentes:

SUPPORTS_CONCURRENT_SESSIONS: bool = True

No aplicativo do servidor, defina o limite de concorrência:

app = create_app(
    create_my_environment,
    MyAction,
    MyObservation,
    max_concurrent_envs=64,  # match or exceed generation_batch_size
)

max_concurrent_envs deve ser ≥ generation_batch_size (que por padrão é per_device_train_batch_size × gradient_accumulation_steps). Por exemplo, com gradient_accumulation_steps=64 e tamanho de lote 1, você precisa de pelo menos 64 sessões simultâneas.

environment_factory vs rollout_func

O GRPOTrainer suporta duas abordagens para treinamento baseado em ambiente:

environment_factory (recomendado): você define uma classe de ambiente com métodos de ferramenta, e o trainer trata da geração, da análise de chamadas de ferramenta e do loop de múltiplas rodadas automaticamente. Essa é a abordagem usada ao longo deste guia.
rollout_func: Você escreve todo o loop de geração e interação com o ambiente por conta própria. Isso dá controle total sobre como as conclusões são produzidas, como as ferramentas são executadas e como as recompensas são calculadas.

Use rollout_func quando environment_factory não se adequar ao seu caso de uso. Por exemplo, servidores de agentes externos em que um servidor externo controla o loop de geração e gerencia seu próprio protocolo de interação entre agente e ambiente.

Treinamento em harnesses: treinando agentes de codificação reais (Harbor)

As integrações acima são white-box: o TRL controla o loop de múltiplas rodadas por conta própria. Ele amostra cada rodada, analisa as chamadas de ferramentas, as executa e devolve os resultados.

Alguns agentes não podem ser conduzidos dessa forma porque possuem seu próprio loop. Um harness de agente de codificação de produção, como opencode ou claude-code, tem seu próprio planejador, conjunto de ferramentas, gerenciamento de contexto e condição de parada. Você deseja treinar exatamente esse agente, e não uma reimplementação dele.

Para isso, o TRL oferece um caminho experimental de caixa-preta (loop-owning) construído sobre experimental.async_grpo.AsyncGRPOTrainer: HarnessAsyncGRPOTrainer controla uma fábrica de sessões do OpenEnv. Os agentes e seus sandboxes vêm do Harbor, servidos através do OpenEnv com openenv harbor serve, de modo que todo agente instalado pelo Harbor é treinável da mesma forma: você escolhe um com --harness. Veja examples/async_grpo_harbor para um exemplo completo, incluindo um launcher do Hugging Face Jobs.

Este caminho requer openenv>=0.7.0, a primeira versão com a API de captura TrainingTrace: pip install "openenv>=0.7.0".

Como funciona

O TRL não amostra cada turno aqui. O agente executa até a conclusão por conta própria, e o TRL lê de volta o que ele fez:

O servidor OpenEnv Harbor inicia o agente em uma tarefa do Harbor em um sandbox (por exemplo, E2B). Cada chamada de modelo feita pelo agente passa pelo proxy de captura do servidor para o seu servidor vLLM, que registra os ids de tokens de prompt e de conclusão do próprio mecanismo, os logprobs amostrados e uma máscara de perda por token.
Quando o agente para, o verificador da própria tarefa pontua o resultado, e o TRL recebe o TrainingTrace validado do OpenEnv e constrói as linhas de treinamento a partir de seus ids de tokens e máscaras. O TRL não retokeniza os prompts capturados.
O GRPO treina nessas linhas. A recompensa é propagada para todos os tokens treinados por meio da vantagem relativa ao grupo, e os novos pesos são sincronizados de volta para o mesmo servidor vLLM.

Cada rollout é executado em sua própria sessão isolada, com seu próprio sandbox e sua própria sessão de captura. O isolamento é importante por dois motivos: o proxy precisa capturar exatamente os tokens daquele rollout, e um rollout não deve interferir em outro. A configuração max_concurrent_rollouts controla quantos rollouts são executados ao mesmo tempo, e o MAX_CONCURRENT_ENVS do servidor deve ser pelo menos max_concurrent_rollouts + 1 (a fábrica mantém uma conexão extra para metadados da tarefa).

Conexões

HarnessAsyncGRPOTrainer recebe o modelo, a configuração e o conjunto de dados da mesma forma que experimental.async_grpo.AsyncGRPOTrainer, com uma fábrica de sessões no lugar de reward_funcs, tools e environment_factory. Sem um harness_adapter, o agente controla seu próprio loop. O trainer constrói a fábrica de sessões com sua própria política de amostragem, então você deve passar um chamável, como partial(HarborSessionFactory, ...), em vez de uma fábrica já construída. Cada sessão retorna um TrainingTrace do OpenEnv contendo os tokens exatos do engine, os logprobs de comportamento e as máscaras pré-preenchidas. rollout_reward_fn é opcional; sem ele, o TRL usa a recompensa do verificador da sessão.

from functools import partial

from harbor_env.harness import HarborSessionFactory

from trl.experimental.async_grpo 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()

O que você precisa definir

O OpenEnv é responsável pela validação da captura e pela seleção de tokens. O TRL converte a máscara de sequência completa do produtor para sua máscara interna de conclusão, empacota as linhas resultantes e aplica o GRPO. Ele não reconstrói tokens nem infere quais chamadas pertencem ao agente.

rollout_reward_fn : como pontuar um rollout

Callable[[HarnessRolloutOutcome], float | None]

O TRL não sabe o que significa sucesso para a sua tarefa, então você converte cada rollout concluído em uma recompensa escalar. A função recebe um HarnessRolloutOutcome que descreve o que o agente fez:

env_reward (float | None): a recompensa do verificador da tarefa, ou None quando o rollout não pôde ser pontuado.
completion (list[dict]): o transcript final da mensagem.
trace (list[TraceEntry]): os registros de captura do OpenEnv selecionados, incluindo turnos com máscara zero.
tool_call_count e tool_failure_count (int): quantas chamadas de ferramenta o agente fez e quantas pareceram falhas.
tool_calls_by_name (dict[str, int]): chamadas por ferramenta, por exemplo {"bash": 3, "edit": 2}.
timed_out (bool): se o agente esgotou seu orçamento de tempo.

Retorne um float ou retorne None para marcar o rollout como não pontuável, de modo que ele seja descartado da linha de base do grupo em vez de ser contado como zero. Se você não passar esta função, o env_reward bruto é usado como está. O exemplo do Harbor mantém a correção do verificador como o sinal principal e adiciona um pequeno bônus de eficiência de ferramentas apenas quando a tarefa é resolvida.

Por que não usar apenas o verificador? O verificador é a responsabilidade do ambiente e responde a uma única pergunta, “quão correto foi o resultado”, o que o mantém limpo e reutilizável para avaliação. A recompensa em que você treina é uma decisão separada, feita no momento do treinamento (binarizar a pontuação, penalizar comportamento degenerado, descartar rollouts não pontuáveis). Ela também precisa de sinais que o verificador nunca vê, já que o verificador inspeciona apenas o resultado final, enquanto rollout_reward_fn também recebe a trajetória (contagens de ferramentas, timed_out, o trace). Por exemplo, um rollout pode passar em alguns testes, mas nunca executar bash; apenas rollout_reward_fn pode ver isso e penalizá-lo.

A elegibilidade de tokens vem do OpenEnv

session.fetch_training_trace() retorna um TrainingTrace validado. Cada turno possui um node_id, prompt_token_ids, completion_token_ids, per_token_logps e um loss_mask obrigatório sobre o prompt e a conclusão. As posições do prompt são zero. Máscaras de conclusão parciais são preservadas, incluindo zeros dentro de uma resposta do assistente.

Por exemplo, o prompt [10, 11], a conclusão [20, 21, 22] e a máscara [0, 0, 1, 0, 1] treinam nos tokens 20 e 22. O token 21 permanece como contexto. Uma conclusão totalmente mascarada é mantida para contexto e contabilização de uso.

O produtor exclui chamadas auxiliares e tentativas descartadas. Não há parâmetros train_turn_fn ou agent_turn_fn nesta API. As integrações existentes devem mover a seleção para suas máscaras OpenEnv. O verificador da tarefa permanece independente da elegibilidade de tokens.

Capturas malformadas falham na fronteira do OpenEnv e interrompem o treinamento. Falhas de transporte permanecem não pontuáveis. Histórias reescritas iniciam linhas separadas, preservando os tokens amostrados e os logprobs. Múltiplas linhas ainda compartilham a recompensa do rollout; esta API não altera o ponderamento do GRPO.

Métricas

O treinador do harness registra as métricas de ferramentas do experimental.async_grpo.AsyncGRPOTrainer com os mesmos nomes, com as seguintes diferenças:

Quando o agente controla seu próprio loop (sem harness_adapter), as falhas são estimadas a partir dos resultados das ferramentas na última solicitação: um resultado é considerado falho quando seu texto contém error, failed, traceback ou exception (sem distinção entre maiúsculas e minúsculas), e resultados com a mesma ferramenta e os mesmos 200 primeiros caracteres são contados uma única vez. tools/latency_s e tools/<name>/latency_s não são registrados, pois o rastro do OpenEnv não possui carimbos de data e hora.
Com um harness_adapter, uma chamada de ferramenta é considerada falha quando o harness relata um erro para ela. O tempo entre uma resposta do modelo e a próxima solicitação é dividido igualmente entre as chamadas de ferramentas daquela rodada e registrado como tools/latency_s, portanto, inclui o próprio overhead do harness. A última rodada não é cronometrada.
tools/unknown_name_count, tools/parallel_calls_mean e tools/loop_exhausted_frac vêm do loop de ferramentas embutido e não são registrados.
O próprio trainer não possui ferramentas, portanto tools/call_frequency e tools/failure_frequency apenas contabilizam os grupos em que pelo menos um rollout chamou uma ferramenta.
rollout/correctness_mean é registrado apenas com uma rollout_reward_fn: o env_reward do verificador, médio sobre os rollouts que possuem um.

Requisitos

O servidor OpenEnv Harbor é o responsável pelo conjunto de dados e pelos sandboxes. Ele precisa da credencial do backend do sandbox (por exemplo, E2B_API_KEY):

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

O servidor vLLM deve expor o tool calling e os ids de tokens reais, e habilitar a sincronização de pesos via NCCL para que o agente sempre acione a política atual. Os rollouts e as atualizações de pesos devem usar o mesmo servidor 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"}'

O treinamento com loop próprio fica em trl.experimental e sua API pode mudar. O exemplo fixa a revisão do OpenEnv com a qual foi testado; consulte seu README para a configuração exata.

Update on GitHub

←Kernels Hub OpenReward→