TRL documentation

Tích hợp OpenEnv để huấn luyện LLM với các môi trường

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Tích hợp OpenEnv để huấn luyện LLM với các môi trường

OpenEnv là một framework mã nguồn mở để định nghĩa, triển khai và tương tác với các môi trường trong học tăng cường (RL) và các luồng công việc agentic. Nó cung cấp các API chuẩn hóa cho việc tương tác với môi trường và hỗ trợ chạy các môi trường dưới dạng máy chủ backend (qua WebSocket hoặc thực thi trong container). Bạn có thể tìm thấy bộ sưu tập các môi trường OpenEnv sẵn sàng sử dụng trên Hugging Face Hub.

Hướng dẫn này bao gồm cách tích hợp OpenEnv với TRL. Để tìm hiểu thêm về chính OpenEnv, hãy xem tài liệu OpenEnv.

Các ví dụ OpenEnv có sẵn: grpo_echo (tối giản), grpo_catch, grpo_wordle, grpo_sudoku, grpo_multi_env, grpo_browsergym, grpo_carla, và async_grpo_harbor (cục bộ và HF Jobs).

Khi nào nên sử dụng môi trường

GRPOTrainer có thể được dùng để huấn luyện các agent. Đối với các tác vụ agentic, nó hỗ trợ hai chế độ: tools, trong đó mô hình có thể gọi các hàm bên ngoài nhưng mỗi lần gọi là không có trạng thái và độc lập, và environments, vốn duy trì trạng thái qua các lượt, cho phép tương tác đa lượt thực sự trong đó hành động của agent định hình các quan sát tiếp theo. Hãy dùng environments khi tính liên tục quan trọng — ví dụ, điều hướng trong một trò chơi, duyệt một trang web, hoặc bất kỳ tác vụ nào mà điều agent nhìn thấy tiếp theo phụ thuộc vào những gì nó đã làm trước đó.

Chọn tích hợp môi trường

OpenEnv là đường dẫn native được tài liệu hóa tại đây. Hai tích hợp khác — OpenReward và Harbor — tuân thủ cùng hợp đồng environment_factory và có thể thay thế cho nhau ở cấp độ TRL. Xem sự so sánh giữa các tích hợp môi trường trong hướng dẫn GRPO để chọn hệ sinh thái phù hợp với tác vụ của bạn.

Cài đặt

Các môi trường OpenEnv được lưu trữ dưới dạng Hugging Face Spaces, đây cũng là các kho Git có thể cài đặt bằng 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"

Điều này cài đặt environment client (ví dụ: EchoEnv) để giao tiếp với máy chủ môi trường từ xa qua WebSocket, cùng với các mô hình hành động/quan sát và tất cả các phụ thuộc cần thiết (bao gồm openenv).

Bạn có thể tìm thấy lệnh cài đặt cho bất kỳ môi trường nào trên trang HF Space của nó. Nhấp vào menu ⋮ (ba chấm) và chọn “Use this Space” để xem các hướng dẫn cài đặt.

Bạn cũng có thể cài đặt gói lõi từ PyPI bằng pip install "openenv>=0.3.1", nhưng lưu ý rằng các phụ thuộc cụ thể cho môi trường có thể cần được cài đặt riêng.

Để phát triển, bạn có thể clone kho OpenEnv và cài đặt cục bộ:

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

Mỗi script môi trường trong TRL bao gồm metadata phụ thuộc inline (PEP 723) để bạn cũng có thể chạy chúng trực tiếp bằng uv:

uv run examples/grpo_echo/grpo_echo.py

Điều này tự động cài đặt gói môi trường cần thiết trong một môi trường ảo được cách ly.

Bắt đầu nhanh

Cách nhanh nhất để hiểu tích hợp là một ví dụ hoàn chỉnh. Script echo.py huấn luyện một mô hình với môi trường Echo, nơi phần thưởng cho các câu trả lời dựa trên độ dài văn bản của chúng:

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

Đó là tất cả. Dưới đây là những gì diễn ra bên trong:

  1. environment_factory=EchoToolEnv: Bộ huấn luyện tạo một instance EchoToolEnv cho mỗi lần sinh (truyền class, không phải instance).
  2. reset() được gọi ở đầu mỗi episode để khởi tạo trạng thái. Trả về một chuỗi quan sát (hoặc None).
  3. Khám phá công cụ: Trình huấn luyện sẽ khám phá tất cả các phương thức công khai trên instance của môi trường (ở đây là echo()) và công khai chúng dưới dạng các công cụ gọi hàm. Mỗi phương thức phải có docstring phù hợp với các đối số có kiểu, để trình huấn luyện sử dụng trong việc xây dựng schema công cụ.
  4. Vòng lặp nhiều lượt: Trình huấn luyện tạo ra một câu trả lời, phân tích các lệnh gọi công cụ, thực thi echo(), thêm kết quả vào và tạo lại, cho đến khi mô hình ngừng gọi công cụ hoặc đạt đến max_completion_length.
  5. Hàm phần thưởng: Đọc env.reward từ mỗi instance môi trường sau khi kết thúc episode (trước khi môi trường được reset).
# 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

Dưới đây là đường cong phần thưởng từ quá trình huấn luyện:

Bạn có thể khám phá thêm các ví dụ OpenEnv sẵn sàng sử dụng trong ghi chú ở đầu trang này.

Cách thức hoạt động của environment_factory

GRPOTrainer của TRL hỗ trợ huấn luyện môi trường tương tác thông qua đối số environment_factory. Khi được cung cấp, trainer sẽ tự động xử lý vòng lặp gọi công cụ nhiều lượt: nó tạo ra các hoàn tất, phân tích các lời gọi công cụ, thực thi chúng đối với môi trường và đưa kết quả trở lại cho mô hình. Tất cả đều không cần mã rollout tùy chỉnh.

Yêu cầu đối với lớp Environment

Lớp môi trường của bạn phải tuân theo các quy tắc sau:

  • __init__(self) (tùy chọn): Nếu được cung cấp, hàm này không được nhận đối số nào. Hãy dùng nó để khởi tạo trạng thái hoặc các client. Nếu cần cấu hình bên ngoài (ví dụ: một URL), hãy lấy nó từ phạm vi bao ngoài hoặc các biến cấp module.
  • reset(self, **kwargs): Được gọi ở đầu mỗi episode. Nhận tất cả các cột của dataset dưới dạng các tham số keyword. Trả về một quan sát dạng chuỗi (hoặc None nếu không có quan sát ban đầu).
  • Phương thức công cụ: Bất kỳ phương thức công khai nào (không bắt đầu bằng _) khác với reset đều được tự động phơi bày dưới dạng một công cụ. Mỗi phương thức công cụ phải có docstring chứa mô tả Args:, vì trainer sử dụng chúng để tạo schema công cụ cho mô hình.

Mẹo cho các lớp môi trường

  • Trạng thái cho phần thưởng: Bạn có thể lưu trữ bất kỳ trạng thái nào trên instance của môi trường (ví dụ: self.reward, self.done, v.v.) và truy cập chúng trong hàm phần thưởng của bạn thông qua tham số environments. Tham khảo hướng dẫn Bắt đầu nhanh để xem ví dụ về mẫu này.
  • Xử lý lỗi: Nếu một phương thức công cụ gây ra ngoại lệ (ví dụ: ValueError("Game over.")), trình huấn luyện sẽ bắt ngoại lệ đó và truyền thông báo lỗi ngược lại cho mô hình dưới dạng phản hồi công cụ. Đây là cách được khuyến nghị để báo hiệu rằng một hành động không hợp lệ hoặc tập đã kết thúc.
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"

Công cụ phải là các phương thức riêng lẻ có tên mô tả và các đối số có kiểu (ví dụ: guess(word: str), move(direction: str)). Chúng tôi không khuyến nghị sử dụng các phương thức chung chung như step(action), vì mô hình cần tên công cụ và mô tả đối số có ý nghĩa để học cách gọi công cụ.

Hàm phần thưởng

Các hàm phần thưởng nhận tham số environments (một danh sách các instance môi trường), vì vậy bạn có thể truy cập bất kỳ trạng thái nào được lưu trữ trong suốt episode:

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

Để biết thêm thông tin về các hàm phần thưởng, hãy xem GRPO - Custom Reward Functions.

Mẹo cho các hàm phần thưởng

Một số điều chúng tôi thấy hữu ích khi làm việc với các môi trường OpenEnv và GRPO:

  • Đảm bảo các nhóm khác nhau. GRPO so sánh các câu trả lời trong cùng một nhóm, do đó thứ hạng tương đối quan trọng hơn các giá trị tuyệt đối, và một nhóm mà mọi câu trả lời đều nhận cùng một phần thưởng sẽ không cung cấp tín hiệu. Phần thưởng nhị phân (1.0 cho thành công, 0.0 trong các trường hợp khác) là rõ ràng khi mô hình thành công đôi khi. Khi mô hình hiếm khi thành công, hãy thêm điểm một phần: ví dụ Sudoku cộng thêm tỷ lệ ô đã điền vào phần thưởng của câu đố đã giải, ví dụ Wordle cộng thêm số ô xanh của lần đoán tốt nhất.
  • Kiểm tra trạng thái cuối cùng, không phải con đường. Khi có thể, hãy để môi trường đánh giá kết quả (ví dụ: “mô hình có giải được câu đố không?”) thay vì kiểm tra xem nó có tuân theo một trình tự hành động cụ thể hay không. Điều này cho phép mô hình tự do khám phá các chiến lược của riêng mình.
  • Thử nghiệm phần thưởng trước khi huấn luyện. Chạy một vài tập một cách thủ công (xem notebook ví dụ Wordle) để xác nhận môi trường trả về các phần thưởng hợp lý. Nếu một mô hình có năng lực không thể đạt điểm cao hơn mức cơ sở ngẫu nhiên, tín hiệu phần thưởng có thể cần được điều chỉnh.

max_completion_length trong các tập đa lượt

Tham số max_completion_length giới hạn tổng số token trên toàn bộ cuộc trò chuyện nhiều lượt (tổng hợp tất cả các lần sinh của mô hình + kết quả công cụ), chứ không chỉ giới hạn cho một lần sinh duy nhất. Đối với các môi trường có nhiều lượt (ví dụ: Sudoku với hàng chục nước đi), bạn có thể cần tăng giá trị này lên:

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

Nếu các episode bị cắt ngắn (mô hình dừng giữa trò chơi), đây có thể là nguyên nhân.

Chạy dài đối với một Space

Một Space có thể đóng các kết nối WebSocket duy trì mở trong nhiều giờ (chúng tôi quan sát thấy điều này xảy ra khoảng mỗi 2 giờ). Các client OpenEnv từ phiên bản 0.11.0 trở về trước sau đó sẽ kết nối lại với một phiên môi trường mới, và nếu một ván game đang diễn ra, lệnh step() tiếp theo sẽ chạy trên một ván game mà client chưa bao giờ reset, mà không có thông báo lỗi nào. Kể từ OpenEnv 0.12.0, các client sẽ ném ra SessionLostError thay vào đó, khiến ván game thất bại. Để tránh việc mất kết nối ảnh hưởng đến quá trình huấn luyện, hãy mở một client mới trong reset(), để mỗi kết nối chỉ kéo dài trong một episode, và đặt MAX_CONCURRENT_ENVS cao hơn một chút so với số lượng ván game được chơi đồng thời, vì các phiên cũ được giải phóng một cách bất đồng bộ:

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

Ví dụ nâng cao: Wordle

Hãy huấn luyện một mô hình để chơi Wordle bằng môi trường TextArena. Điều này minh họa tương tác nhiều lượt, xử lý phản hồi tích lũy, phần thưởng được chia thành nhiều thành phần, và kết thúc episode thông qua ngoại lệ.

Hướng dẫn Wordle của OpenEnv bao phủ toàn bộ quy trình: dữ liệu SFT, bước khởi động SFT, bước GRPO này và đánh giá, kèm theo phần cứng và chi phí. Phiên bản notebook của ví dụ này là grpo_wordle.ipynb.

Môi trường TextArena

TextArena là một bộ sưu tập mã nguồn mở gồm các trò chơi văn bản cạnh tranh, được thiết kế để đánh giá kỹ năng suy luận của LLM thông qua các trò chơi văn bản như Wordle, Snake, Tic-Tac-Toe và nhiều trò chơi khác.

image of TextArena

Tại sao là Wordle?

Wordle là một benchmark tốt cho RL dựa trên môi trường vì nó đòi hỏi suy luận về phản hồi, hoàn toàn dựa trên văn bản, và các mô hình từ 1B tham số có thể cải thiện hiệu suất trên nó.

Wordle hoạt động như thế nào? Wordle là một trò chơi đoán từ, trong đó người chơi phải đoán một từ gồm 5 chữ cái trong 6 lượt. Sau mỗi lần đoán, môi trường cung cấp phản hồi theo từng chữ cái:

G U E S S
X G Y X X

X = không có trong từ, G = đúng vị trí (xanh lá), Y = sai vị trí (vàng). Ở đây, “U” là đúng và đúng chỗ, “E” có trong từ nhưng sai chỗ.

Bắt đầu từ bước khởi động SFT

Qwen3-1.7B gần như không bao giờ thắng một ván nào, nên GRPO từ mô hình gốc có rất ít điều để học: một nhóm các ván mà không ván nào được thắng sẽ không tạo ra tín hiệu. Ví dụ bắt đầu từ openenv/Qwen3-1.7B-Wordle-SFT, được huấn luyện trong một epoch trên 1.000 ván do một bộ giải quyết chơi ở cùng định dạng gọi công cụ, với một phần suy luận ngắn trước mỗi lần đoán.

Lớp Environment

Lớp WordleEnv bọc client TextArena, phơi bày guess() dưới dạng công cụ và giữ lại những gì phần thưởng cần:

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

Các lựa chọn thiết kế chính:

  • reset() trả về thông điệp trò chơi ban đầu như quan sát đầu tiên mà mô hình nhìn thấy.
  • guess() là công cụ duy nhất. Mô hình gọi nó mỗi lượt với một từ gồm 5 chữ cái.
  • Cắt lát phản hồi tích lũy: TextArena trả về toàn bộ lịch sử trò chơi ở mỗi lượt. Chúng tôi chỉ cắt ra phần mới để tránh lặp lại ngữ cảnh.
  • Lỗi ngoại lệ khi hoàn tất: Nếu mô hình cố đoán sau khi trò chơi kết thúc, guess() sẽ ném ra ValueError. Trainer sẽ bắt ngoại lệ này và trả về "Game over." cho mô hình như một phản hồi công cụ, đồng thời tính vào tools/failure_frequency.

Hàm phần thưởng và huấn luyện

Phần thưởng của chính môi trường trả một phần thưởng chiến thắng như nhau dù chỉ đoán 1 lần hay 6 lần, và chỉ tính điểm một phần ở cuối trò chơi. Khi bật chế độ suy nghĩ, một chính sách có thể học cách suy nghĩ ít hơn và sử dụng cả 6 lần đoán để đi đến cuối trò chơi. Ba hạng mục phần thưởng giúp tránh điều đó:

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

Chiến thắng càng đến sớm thì phần thưởng càng cao (1.42 nếu đoán đúng trong một lần, 1.0 nếu trong sáu lần), thất bại được tính phần thưởng dựa trên lần đoán tốt nhất của ván đấu bất kể ván có kết thúc hay không, và mỗi nước đi không hợp lệ bị trừ 0.1. TRL ghi log từng thành phần riêng biệt, do đó rewards/reward_win/mean tăng lên cùng với tỷ lệ chiến thắng.

Chạy ví dụ

Mỗi bước chơi 64 ván, mỗi ván là một phiên môi trường. Tạo bản sao của Space openenv/wordle và đặt biến MAX_CONCURRENT_ENVS=64 trên bản sao của bạn.

colocate
server

Chế độ Colocate (1 GPU, được khuyến nghị)

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

Điều này chạy vLLM trong cùng tiến trình với huấn luyện, chỉ yêu cầu một GPU duy nhất. Trong lần chạy của chúng tôi, 150 bước mất khoảng 2 giờ 15 phút trên một A100 80GB.

Kết quả

Từ lần chạy của chúng tôi, với một lần chạy huấn luyện cho mỗi mô hình, trên cùng 30 từ được chơi theo cách tham lam:

Mô hìnhChiến thắngSố lần đoán mỗi lần thắngCác lần đoán không hợp lệ
Qwen3-1.7B0/30-24
+ SFT10/303.241
+ SFT + GRPO (mô hình)20/304.13

GRPO phần lớn đã loại bỏ các từ lặp lại và từ bịa đặt, đồng thời giành được nhiều trận thắng hơn. Phần lớn các trận thắng mới cần từ 4 đến 6 lượt đoán, do đó số lượt đoán trên mỗi trận thắng đã tăng lên. 30 từ này được lấy từ 156 danh từ mà TextArena dùng để rút bí mật, vốn các mô hình cũng đã thấy trong quá trình huấn luyện, và một lần chạy mới sẽ cho kết quả chênh lệch vài trận thắng (hai lần chạy SFT với cùng công thức đã thắng 10 và 13 trận). Xem ba mô hình chơi trong Wordle Arena và xem các đường cong huấn luyện trên Trackio.

Đào tạo đa môi trường

Bạn có thể huấn luyện một mô hình duy nhất trên nhiều môi trường cùng lúc. Điều này hữu ích khi bạn muốn mô hình học các kỹ năng khác nhau song song. Ví dụ, chơi Wordle (lý luận ngôn ngữ) và Catch (lý luận không gian) trong cùng một lượt huấn luyện.

Ý tưởng chính là tạo ra một lớp môi trường siêu (meta-environment class) bao bọc nhiều môi trường và định tuyến mỗi mẫu đến môi trường tương ứng bằng cách sử dụng một cột trong bộ dữ liệu.

Cách hoạt động

  1. Thêm một cột "env" (hoặc tương tự) vào bộ dữ liệu của bạn để xác định môi trường mà mỗi mẫu thuộc về.
  2. Trong reset(**kwargs), đọc kwargs["env"] để chọn môi trường hoạt động cho tập đó.
  3. Tiết lộ các công cụ từ tất cả các môi trường; trình huấn luyện sẽ phát hiện tất cả các phương thức công khai.
  4. Sử dụng các hàm phần thưởng riêng cho từng môi trường, trả về None cho các mẫu không thuộc môi trường đó. TRL xử lý các giá trị None bằng nansum/nanmean.

Ví dụ: Wordle + Catch

Script grpo_multi_env.py huấn luyện đồng thời trên Wordle và 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. ..."""
        ...

Các mẫu chính:

  • Khởi tạo client theo kiểu lười: Tạo các client trong reset(), không phải trong __init__(), để tránh các kết nối WebSocket không cần thiết.
  • Đóng trước khi mở lại: Đóng client trước đó trước khi tạo một client mới để tránh lỗi dung lượng máy chủ.
  • kwargs routing: Cột "env" từ dataset được truyền vào reset() dưới dạng tham số keyword.
  • Mọi công cụ đều được hiển thị đồng thời: Mô hình nhìn thấy guess, move và stay như các công cụ khả dụng bất kể môi trường đang hoạt động là gì. Nếu nó gọi sai công cụ (ví dụ: move trong Wordle), phương pháp sẽ ném ra ValueError mà trình huấn luyện sẽ bắt một cách mềm mại. Trên thực tế, các mô hình học cách sử dụng đúng công cụ dựa trên system prompt.

Hàm phần thưởng cho từng môi trường

Mỗi hàm phần thưởng trả về None cho các mẫu từ các môi trường khác:

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 chuyển đổi None thành nan bên trong và sử dụng nansum/nanmean để tổng hợp, do đó mỗi mẫu chỉ được chấm điểm bởi hàm phần thưởng liên quan của nó.

Bộ dữ liệu với định tuyến môi trường

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

Chạy ví dụ đa môi trường

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

Khi huấn luyện trên nhiều môi trường, hãy theo dõi các chỉ số theo từng hàm phần thưởng (train/reward_func_0, train/reward_func_1, v.v.) thay vì chỉ số tổng hợp train/reward. Chỉ số tổng hợp luân phiên giữa các môi trường và có thể trông nhiễu.

Chạy các môi trường

Khi sử dụng environment_factory, trainer sẽ tự động kết nối với máy chủ môi trường. Bạn chỉ cần đảm bảo máy chủ đang chạy. Có ba cách để chạy một máy chủ môi trường OpenEnv:

space
docker
local

Kết nối với một Hugging Face Space từ xa (đơn giản nhất)

Hầu hết các script ví dụ mặc định sử dụng một Space được lưu trữ (không cần thiết lập):

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

Để huấn luyện, hãy nhân bản Space sang tài khoản của bạn để tránh các vấn đề về đồng thời. Trình huấn luyện mở N kết nối WebSocket đồng thời (một kết nối cho mỗi lần sinh), và các Space dùng chung có thể không hỗ trợ điều này. Xem Server concurrency để biết thêm chi tiết.

Danh mục môi trường

Cách tốt nhất để khám phá danh mục các môi trường đang được bảo trì hiện tại là truy cập danh mục chính thức của OpenEnv.

Để tạo môi trường của riêng bạn, hãy xem hướng dẫn về Xây dựng Môi trường của riêng bạn với OpenEnv. Các môi trường được tích hợp chặt chẽ với Hub, vì vậy bạn có thể đẩy các môi trường mới để cộng đồng có thể tái sử dụng.

Độ đồng thời của máy chủ

Khi sử dụng environment_factory, trainer tạo ra N instance môi trường (một instance cho mỗi lần generation), mỗi instance mở một kết nối WebSocket đến server. Mặc định, các server OpenEnv chỉ cho phép 1 phiên đồng thời, điều này sẽ gây ra lỗi trong quá trình huấn luyện.

Để hỗ trợ huấn luyện song song, hãy cấu hình máy chủ cho tính đồng thời:

  1. Trong tệp môi trường của bạn, khai báo hỗ trợ phiên đồng thời:
SUPPORTS_CONCURRENT_SESSIONS: bool = True
  1. Trong ứng dụng máy chủ của bạn, hãy đặt giới hạn đồng thời:
app = create_app(
    create_my_environment,
    MyAction,
    MyObservation,
    max_concurrent_envs=64,  # match or exceed generation_batch_size
)

max_concurrent_envs phải ≥ generation_batch_size (mặc định là per_device_train_batch_size × gradient_accumulation_steps). Ví dụ, với gradient_accumulation_steps=64 và batch size là 1, bạn cần ít nhất 64 phiên đồng thời.

environment_factory so với rollout_func

GRPOTrainer hỗ trợ hai phương pháp cho việc huấn luyện dựa trên môi trường:

  • environment_factory (khuyến nghị): Bạn định nghĩa một lớp môi trường với các phương thức công cụ, và trình huấn luyện sẽ tự động xử lý việc tạo dữ liệu, phân tích lời gọi công cụ và vòng lặp đa lượt. Đây là cách tiếp cận được sử dụng xuyên suốt hướng dẫn này.
  • rollout_func: Bạn tự viết toàn bộ vòng lặp sinh và tương tác với môi trường. Điều này cho phép kiểm soát hoàn toàn cách các câu trả lời được tạo ra, cách các công cụ được thực thi và cách phần thưởng được tính toán.

Sử dụng rollout_func khi environment_factory không phù hợp với trường hợp sử dụng của bạn. Ví dụ, các máy chủ tác nhân bên ngoài trong đó một máy chủ bên ngoài sở hữu vòng lặp sinh và tự quản lý giao thức tương tác tác nhân-môi trường của riêng nó.

Huấn luyện trên các bộ khung: huấn luyện các tác nhân lập trình thực sự (Harbor)

Các tích hợp ở trên là white-box: TRL tự điều khiển vòng lặp đa lượt. Nó lấy mẫu mỗi lượt, phân tích các lệnh gọi công cụ, thực thi chúng và đưa kết quả trở lại.

Một số agent không thể được điều khiển theo cách này vì chúng tự có vòng lặp riêng. Một harness agent lập trình production như opencode hoặc claude-code có bộ lập kế hoạch, bộ công cụ, quản lý ngữ cảnh và điều kiện dừng riêng. Bạn muốn huấn luyện chính xác agent đó, chứ không phải một bản tái hiện lại của nó.

Để làm điều này, TRL cung cấp một đường dẫn thử nghiệm black-box (loop-owning) được xây dựng trên AsyncGRPOTrainer: HarnessAsyncGRPOTrainer điều khiển một nhà máy phiên OpenEnv. Các tác nhân và sandbox của chúng đến từ Harbor, được phục vụ thông qua OpenEnv bằng openenv harbor serve, do đó mọi tác nhân mà Harbor cài đặt đều có thể được huấn luyện theo cùng một cách: bạn chọn một tác nhân bằng --harness. Xem examples/async_grpo_harbor để có một ví dụ hoàn chỉnh, bao gồm cả trình khởi chạy Hugging Face Jobs.

Lộ trình này yêu cầu openenv>=0.7.0, phiên bản đầu tiên có API bắt TrainingTrace: pip install "openenv>=0.7.0".

Cách hoạt động

TRL không lấy mẫu từng lượt ở đây. Tác nhân tự chạy đến khi hoàn tất, và TRL đọc lại những gì nó đã thực hiện:

  1. Máy chủ OpenEnv Harbor khởi động tác nhân trên một tác vụ Harbor trong một môi trường sandbox (ví dụ E2B). Mỗi lần gọi mô hình mà tác nhân thực hiện đều đi qua proxy ghi nhận của máy chủ đến máy chủ vLLM của bạn, nơi ghi lại các id token prompt và completion của chính engine, các logprobs được lấy mẫu, và một loss mask cho từng token.
  2. Khi agent dừng lại, bộ xác thực riêng của tác vụ sẽ chấm điểm kết quả, và TRL nhận TrainingTrace đã được xác thực từ OpenEnv rồi xây dựng các hàng huấn luyện từ các token id và masks của nó. TRL không thực hiện lại quá trình token hóa cho các prompt đã được ghi nhận.
  3. GRPO huấn luyện trên các dòng đó. Phần thưởng được lan truyền đến mọi token được huấn luyện thông qua lợi thế tương đối theo nhóm, và trọng số mới được đồng bộ hóa ngược lại vào cùng một máy chủ vLLM.

Mỗi rollout được chạy trong phiên cách ly riêng, với sandbox riêng và phiên capture riêng. Sự cách ly này quan trọng vì hai lý do: proxy phải capture chính xác các token của rollout đó, và một rollout không được phép gây ảnh hưởng đến rollout khác. Cài đặt max_concurrent_rollouts kiểm soát số lượng rollout chạy đồng thời, và MAX_CONCURRENT_ENVS của server phải ít nhất bằng max_concurrent_rollouts + 1 (factory giữ lại một kết nối bổ sung cho metadata của task).

Kết nối

HarnessAsyncGRPOTrainer nhận model, config và dataset giống như AsyncGRPOTrainer, với một session factory thay thế cho reward_funcs, tools và environment_factory. Nếu không có harness_adapter, agent sẽ tự quản lý vòng lặp của mình. Trainer xây dựng session factory với chính sách sampling của riêng nó, vì vậy bạn truyền vào một callable như partial(HarborSessionFactory, ...) thay vì một factory đã được xây dựng sẵn. Mỗi session trả về một OpenEnv TrainingTrace chứa chính xác các token của engine, behavior logprobs và các mask đã được prefill. rollout_reward_fn là tùy chọn; nếu không có nó, TRL sẽ sử dụng reward từ verifier của session.

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

Những gì bạn cần định nghĩa

OpenEnv chịu trách nhiệm xác thực dữ liệu thu thập và lựa chọn token. TRL chuyển đổi mặt nạ chuỗi đầy đủ của bên tạo ra thành mặt nạ hoàn thành nội bộ, đóng gói các hàng kết quả và áp dụng GRPO. Nó không tái tạo token hay suy luận xem những cuộc gọi nào thuộc về tác tử.

rollout_reward_fn : cách chấm điểm một lượt rollout

Callable[[HarnessRolloutOutcome], float | None]

TRL không biết thành công đối với tác vụ của bạn là gì, vì vậy bạn chuyển mỗi rollout hoàn thành thành một phần thưởng vô hướng. Hàm này nhận vào một HarnessRolloutOutcome mô tả những gì tác nhân đã thực hiện:

  • env_reward (float | None): phần thưởng từ bộ xác minh của tác vụ, hoặc None khi rollout không thể được chấm điểm.
  • completion (list[dict]): bản ghi tin nhắn cuối cùng.
  • trace (list[TraceEntry]): các bản ghi capture OpenEnv đã được chọn, bao gồm các lượt bị zero-mask.
  • tool_call_count và tool_failure_count (int): số lần agent gọi công cụ và số lần có vẻ như thất bại.
  • tool_calls_by_name (dict[str, int]): số lần gọi cho từng công cụ, ví dụ {"bash": 3, "edit": 2}.
  • timed_out (bool): cho biết agent đã hết ngân sách thời gian hay chưa.

Trả về một float, hoặc trả về None để đánh dấu rollout không thể chấm điểm, giúp nó bị loại khỏi baseline của nhóm thay vì được tính là số không. Nếu bạn không truyền hàm này, env_reward thô sẽ được sử dụng nguyên vẹn. Ví dụ Harbor giữ độ chính xác của verifier làm tín hiệu chính và chỉ thêm một phần thưởng nhỏ về hiệu quả công cụ khi tác vụ được giải quyết.

Vì sao không chỉ dùng verifier? Verifier là trách nhiệm của môi trường và chỉ trả lời một câu hỏi: “kết quả đúng đến mức nào”, nhờ đó nó giữ được sự sạch sẽ và có thể tái sử dụng cho việc đánh giá. Phần thưởng (reward) mà bạn dùng để huấn luyện là một quyết định riêng biệt, chỉ áp dụng trong thời gian huấn luyện (binarize điểm, phạt hành vi suy thoái, loại bỏ các rollout không chấm điểm được). Nó cũng cần các tín hiệu mà verifier không bao giờ nhìn thấy, vì verifier chỉ kiểm tra kết quả cuối cùng, trong khi rollout_reward_fn còn nhận được cả trajectory (số lượng tool, timed_out, trace). Ví dụ, một rollout có thể vượt qua một số bài test nhưng không bao giờ chạy bash; chỉ rollout_reward_fn mới nhìn thấy điều đó và có thể phạt nó.

Tính hợp lệ của token đến từ OpenEnv

session.fetch_training_trace() trả về một TrainingTrace đã được xác thực. Mỗi lượt có một node_id, prompt_token_ids, completion_token_ids, per_token_logps duy nhất và một loss_mask bắt buộc trên cả prompt lẫn completion. Các vị trí prompt có giá trị bằng không. Các mask completion một phần được giữ nguyên, bao gồm cả các giá trị không bên trong phản hồi của assistant.

Ví dụ, prompt [10, 11], completion [20, 21, 22] và mask [0, 0, 1, 0, 1] sẽ huấn luyện trên các token 20 và 22. Token 21 vẫn được giữ làm ngữ cảnh. Một completion bị mask hoàn toàn sẽ được giữ lại cho ngữ cảnh và để tính toán mức sử dụng.

Producer loại trừ các lời gọi phụ và các lần thử lại bị hủy bỏ. API này không có các tham số train_turn_fn hoặc agent_turn_fn. Các tích hợp hiện có phải di chuyển việc lựa chọn vào các OpenEnv masks của chúng. Bộ xác minh của tác vụ vẫn độc lập với tính đủ điều kiện của token.

Các bản ghi không hợp lệ sẽ thất bại tại ranh giới OpenEnv và làm dừng quá trình huấn luyện. Các lỗi truyền tải vẫn không thể chấm điểm được. Các lịch sử được viết lại sẽ bắt đầu các hàng riêng biệt, giữ nguyên các token và logprobs đã lấy mẫu. Nhiều hàng vẫn chia sẻ phần thưởng rollout; API này không thay đổi trọng số của GRPO.

Chỉ số

Trainer của harness ghi lại chỉ số công cụ của AsyncGRPOTrainer dưới cùng các tên, với những khác biệt sau:

  • Khi agent tự quản lý vòng lặp của mình (không có harness_adapter), các lỗi được ước tính từ kết quả công cụ trong yêu cầu gần nhất: một kết quả được tính là thất bại khi văn bản của nó chứa error, failed, traceback hoặc exception (không phân biệt hoa thường), và các kết quả có cùng công cụ cùng 200 ký tự đầu tiên chỉ được tính một lần. tools/latency_s và tools/<name>/latency_s không được ghi log, vì trace của OpenEnv không có dấu thời gian.
  • Với harness_adapter, một lời gọi công cụ được tính là thất bại khi harness báo lỗi cho nó. Khoảng thời gian giữa một phản hồi của mô hình và yêu cầu tiếp theo được chia đều cho các lời gọi công cụ của lượt đó và được ghi nhận là tools/latency_s của chúng, do đó bao gồm cả chi phí xử lý của chính harness. Lượt cuối cùng không được đo thời gian.
  • tools/unknown_name_count, tools/parallel_calls_mean và tools/loop_exhausted_frac đến từ vòng lặp công cụ tích hợp sẵn và không được ghi log.
  • Bản thân trainer không có công cụ nào, do đó tools/call_frequency và tools/failure_frequency chỉ đếm các nhóm mà trong đó ít nhất một rollout đã gọi công cụ.
  • rollout/correctness_mean chỉ được ghi log khi có rollout_reward_fn: đó là env_reward của bộ xác minh, được lấy trung bình trên các rollout có giá trị này.

Yêu cầu

Máy chủ OpenEnv Harbor sở hữu bộ dữ liệu và các sandbox. Nó cần thông tin xác thực của backend sandbox (ví dụ E2B_API_KEY):

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

Máy chủ vLLM phải hỗ trợ gọi công cụ và các id token thực, đồng thời bật đồng bộ trọng số NCCL để tác nhân luôn truy cập chính sách hiện tại. Các lượt rollout và cập nhật trọng số phải sử dụng cùng một máy chủ 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"}'

Huấn luyện dạng loop-owning nằm trong trl.experimental và API của nó có thể thay đổi. Ví dụ này cố định phiên bản OpenEnv mà nó đã được kiểm thử; xem README của nó để biết cách thiết lập chính xác.

Update on GitHub