Integração do OpenReward para Treinar LLMs com Ambientes

OpenReward é um ecossistema aberto para ambientes de RL construído sobre o Open Reward Standard (ORS) — um protocolo público e agnóstico de linguagem HTTP/SSE para a forma como um ambiente expõe suas tarefas, ferramentas, sessões e recompensas. Como o ORS é apenas um protocolo, o mesmo ambiente pode ser executado na OpenReward platform, auto-hospedado em qualquer serviço de contêineres ou localmente em localhost para desenvolvimento. Um catálogo de ambientes prontos para uso está disponível em openreward.ai.

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

A integração está em trl.experimental.openreward e é controlada pelo extra trl[openreward] (importação preguiçosa — não usuários não pagam nada).

Quando usar ambientes OpenReward

GRPOTrainer suporta treinamento baseado em ambiente por meio do slot environment_factory — consulte OpenEnv para o contrato geral. Use OpenReward quando quiser treinar contra um ambiente que fala ORS: o catálogo OpenReward (por exemplo, Eigent/SETA, kanishk/EndlessTerminals, nebius/SWE-rebench-V2), um ambiente que você hospeda na sua própria infraestrutura ou um servidor local que você está desenvolvendo.

Instalação

pip install trl[openreward]

Isso instala o SDK Python openreward. A própria integração importa openreward de forma preguiçosa, de modo que usuários que não interagem com trl.experimental.openreward não são afetados.

Início rápido

A classe OpenRewardSpec conecta um único ambiente ORS aos três slots do treinador TRL — train_dataset, environment_factory, reward_funcs — expondo propriedades que mapeiam 1:1 para esses nomes de kwarg:

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

Por baixo dos panos, OpenRewardSpec faz três coisas, de forma preguiçosa no primeiro acesso:

1. spec.train_dataset: deriva um datasets.Dataset a partir da lista de tarefas do env (uma única ida e volta HTTP via o SDK). Possui no mínimo prompt, task_index, além de colunas de metadados por tarefa incorporadas.
2. spec.environment_factory: retorna um chamável sem argumentos que produz um novo adaptador por rollout a cada chamada. O adaptador expõe um método Python por ferramenta ORS, com assinatura tipada e docstring geradas automaticamente a partir do JSON Schema do env. O coletor de ferramentas do TRL as identifica por meio de inspect.getmembers.
3. spec.reward_funcs: uma função de recompensa apenas de resultado (última recompensa não nula na trajetória) adequada para ambientes de recompensa esparsa, como o SETA.

Usando um ambiente de hub

Passe um nome de catálogo do openreward.ai como destino. O SDK lê OPENREWARD_API_KEY do ambiente para autenticação.

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

Usando um ambiente auto-hospedado

Passe a URL diretamente. Nenhuma chave de API é necessária se o seu servidor não exigir uma.

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

O SDK openreward espera, por padrão, um layout de plataforma com dois subdomínios (api.<host> para chamadas stateless e sessions.<host> para chamadas de sessão baseadas em SSE). Para servidores self-hosted de host único (uma única URL servindo tudo), defina as variáveis de ambiente de override abaixo antes de construir o OpenRewardSpec:

Copiado

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

Executando um ambiente mínimo localmente

A maneira mais rápida de testar a integração de ponta a ponta sem dependências externas é um pequeno servidor ORS definido com o andaime Environment + Server do SDK openreward. O exemplo abaixo é um ambiente echo completo — o modelo vence ao chamar echo(text=...) com a string alvo da tarefa.

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

Execute-o:

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

Em seguida, aponte OpenRewardSpec para ele (com as substituições de URL descritas acima):

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

Esse também é o padrão de fixture usado pelos próprios testes do TRL — veja trl-internal-testing/openreward-echo-env para o Space implantado.

Selecionando tarefas

OpenRewardSpec aceita uma contagem ou uma lista de índices explícita:

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 e indices são mutuamente exclusivos e ambos buscam apenas as tarefas de que precisam (sem varredura da lista completa de tarefas).

Como o vínculo de ferramentas funciona

Na construção, a spec chama o endpoint /tools do env para buscar uma lista de specs de ferramentas (cada uma com um nome, uma descrição e um JSON Schema para os argumentos). Para cada ferramenta, ele gera um método Python no adapter por rollout com uma assinatura tipada e um docstring derivado do schema. Os métodos são definidos na classe do adapter, onde o coletor de ferramentas do TRL os encontra, de modo que transformers.utils.get_json_schema e inspect.getmembers(type(env), isfunction) do TRL produzem o schema de ferramenta correto para o modelo sem código de wrapper por env.

Se a descrição de uma ferramenta contiver caracteres que não são seguros para serem inseridos em código-fonte Python, o binder recorre a uma forma sanitizada para que a vinculação nunca falhe em ambientes reais.

Funções de recompensa

spec.reward_funcs tem como padrão uma recompensa apenas de resultado — para cada rollout, retorna a última recompensa não nula observada durante a trajetória. Esse é o padrão adequado para ambientes de recompensa esparsa (por exemplo, SETA, onde apenas submit_solution retorna uma recompensa não nula).

Se nenhuma ferramenta retornar uma recompensa não nula, a função de recompensa padrão retorna None, e não 0.0. Isso inclui rollouts sem chamadas de ferramentas, apenas com chamadas de ferramentas não pontuadas ou apenas com chamadas de ferramentas falhas. Uma recompensa observada de 0.0 ou negativa permanece um score válido. Uma chamada de ferramenta posterior não pontuada ou falha não apaga um score anterior.

O GRPO exclui dos baselines de recompensa os rollouts para os quais todas as funções de recompensa retornam None e lhes atribui vantagem zero. Isso não remove seus tokens da normalização da perda nem desativa outros termos, como a regularização KL. Se outra função de recompensa pontuar o rollout, ele ainda pode receber uma vantagem não nula.

Pontuações ausentes não identificam por que um episódio terminou. Se desistir ou atingir o limite de chamadas de ferramentas deve contar como uma falha na tarefa, retorne uma recompensa de falha explícita do seu ambiente ou função de recompensa personalizada. Falhas de infraestrutura ou do avaliador podem, em vez disso, permanecer sem pontuação; monitore a cobertura de pontuação junto com a recompensa média.

Se você quiser uma recompensa personalizada, escreva uma função de recompensa TRL regular e passe-a diretamente:

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

O adaptador por rollout expõe o estado em execução que o TRL precisa — env.reward, env.rewards, env.metadata, env.finished, env.last_output — para modelagem pós-hoc de recompensas arbitrária.

OpenRewardSpec

Limitações

• A integração está em trl.experimental — as APIs podem mudar. Defina TRL_EXPERIMENTAL_SILENCE=1 para silenciar o aviso nos logs de CI.
• Atualmente, expõe um único OpenRewardSpec que cobre um único ambiente; o treinamento em múltiplos ambientes (no padrão de “meta-ambiente” do OpenEnv) ainda não é suportado.
• Rollouts de longa duração (>15 min por episódio) precisam de um ping de keepalive — ainda não implementado.

Referência

• Padrão de Recompensa Aberto
• Plataforma OpenReward
• SDK Python openreward
• Echo env Space — trl-internal-testing/openreward-echo-env