Repository navigation
Add experimental Harbor integration for GRPO environment training #6018
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 4 commits
b5c0768
1f6f01c
4d3252d
bb446bb
001cd6e
a76c804
0679c5c
ca56638
ba28a18
7600e96
08aae44
0dafb6a
f73e92f
adcf95d
47fff48
c39cd77
678bbc9
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,159 @@ | ||
| # Harbor Integration for Training LLMs with Environments | ||
|
|
||
| [Harbor](https://www.harborframework.com) is a framework for running agentic tasks in sandboxes. It decouples a **task** (instruction + sandbox image + verifier), a **harness/agent** (the tool surface + loop), and a **sandbox** (`docker`, `e2b`, `daytona`, `gke`, …) so they can be mixed freely. This makes it a natural fit for RL: the same task suite can be trained with different tool surfaces, on whichever sandbox backend you prefer. | ||
|
|
||
| This guide covers **how to integrate Harbor with TRL**. For Harbor itself, see the [Harbor docs](https://www.harborframework.com/docs). | ||
|
|
||
| > [!NOTE] | ||
| > The integration lives at `trl.experimental.harbor` and is gated behind the `trl[harbor]` extra (lazy-imported — non-users pay nothing). | ||
|
|
||
| ## When to use Harbor environments | ||
|
|
||
| [`GRPOTrainer`] supports environment-based training via the `environment_factory` slot — see [OpenEnv](openenv) for the general contract. Use Harbor when you want to train against a **Harbor task suite**: a directory tree of tasks, each a self-contained sandbox + verifier (for example, a data-analysis agent suite where the model explores files in a sandbox and writes an answer that a grader checks). | ||
|
|
||
| ## Installation | ||
|
|
||
| ```bash | ||
| pip install trl[harbor] | ||
| ``` | ||
|
|
||
| This installs the `harbor` framework (Python >= 3.12). The integration imports `harbor` lazily and runs it **in-process**, so users who don't touch `trl.experimental.harbor` aren't affected. | ||
|
|
||
| A sandbox backend must also be installed and reachable at train time. Harbor keeps cloud backends behind its own extras, so install the one you intend to use and provide its credentials: | ||
|
|
||
| ```bash | ||
| pip install "harbor[e2b]" # E2B cloud sandbox -> environment_type="e2b", needs E2B_API_KEY | ||
| # docker backend (environment_type="docker", Harbor's default) just needs a reachable Docker daemon | ||
| ``` | ||
|
|
||
| ## Quick start | ||
|
|
||
| `HarborSpec` wires a single Harbor task suite into the three TRL trainer slots — `train_dataset`, `environment_factory`, `reward_funcs` — by exposing properties that map 1:1 to those kwarg names: | ||
|
|
||
| ```python | ||
| from trl import GRPOConfig, GRPOTrainer | ||
| from trl.experimental.harbor import HarborSpec | ||
|
|
||
| spec = HarborSpec("AdithyaSK/data_agent_rl_environment_train", agent="bash", num_tasks=64) | ||
|
adithya-s-k marked this conversation as resolved.
Resolved
|
||
|
|
||
| trainer = GRPOTrainer( | ||
| model="Qwen/Qwen3-4B", | ||
| args=GRPOConfig( | ||
| num_generations=8, | ||
| max_steps=50, | ||
| max_tool_calling_iterations=25, | ||
| log_completions=True, | ||
| ), | ||
| train_dataset=spec.train_dataset, | ||
| environment_factory=spec.environment_factory, | ||
| reward_funcs=spec.reward_funcs, | ||
| ) | ||
| trainer.train() | ||
| ``` | ||
|
|
||
| Under the hood `HarborSpec` does three things, lazily on first access: | ||
|
|
||
| 1. **`spec.train_dataset`**: resolves the task suite to local task directories (downloading the HF dataset if needed) and builds a `datasets.Dataset` with `prompt` (empty — the env's instruction is appended at `reset`), `task_dir`, `task_index`, plus per-task `task.toml` metadata columns. | ||
| 2. **`spec.environment_factory`**: returns a zero-arg callable producing a fresh per-rollout [`~trl.experimental.harbor.HarborEnv`]. On `reset(task_dir)` it starts the task's Harbor sandbox and returns its instruction; tool methods exec into the sandbox; `env.reward` runs the verifier once after the rollout. | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fixed — all references now use the
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. these aren't links yet for me. no?
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good catch — the cross-refs only resolve when the target has an autodoc entry. Added |
||
| 3. **`spec.reward_funcs`**: an outcome reward that reads the Harbor verifier's scalar per rollout. | ||
|
|
||
| ## The dataset | ||
|
|
||
| `dataset` is either a Hugging Face dataset repo id holding a Harbor task tree, or a local path containing a `tasks/` subtree. Each task is a directory: | ||
|
|
||
| ``` | ||
| tasks/<task_id>/ | ||
| ├── instruction.md # the task prompt (returned by reset) | ||
| ├── task.toml # config + metadata (gold answer, difficulty, ...) | ||
| ├── environment/ # Dockerfile (+ any pre-agent data hooks) | ||
| └── tests/ # test.sh / grader → writes the reward | ||
| ``` | ||
|
|
||
| Select a subset with `num_tasks` or `indices` (mutually exclusive): | ||
|
|
||
| ```python | ||
| spec = HarborSpec("AdithyaSK/data_agent_rl_environment_train", num_tasks=10) # first 10 | ||
| spec = HarborSpec("AdithyaSK/data_agent_rl_environment_train", indices=[0, 5, 13]) # specific | ||
| ``` | ||
|
|
||
| ## Agents: external vs installed | ||
|
|
||
| Harbor supports two ways an agent drives a task, and the distinction determines what can be trained with RL: | ||
|
|
||
| - [**External agents**](https://www.harborframework.com/docs/agents#external-agents) run *outside* the sandbox and drive the loop themselves, issuing commands into the container through Harbor's environment interface ("typically by executing bash commands via the `exec` method"). The agent decides each action and interprets each result; the sandbox only executes. | ||
| - [**Installed agents**](https://www.harborframework.com/docs/agents#installed-agents) are installed *into the container image* and run there as a headless subprocess (extending `BaseInstalledAgent`). Harbor launches the agent inside the sandbox and parses its trajectory file afterward (`populate_context_post_run`); the agent runs autonomously with its own inference. | ||
|
|
||
| **TRL's integration is the external-agent pattern, and only that pattern is supported for now.** RL training requires the trainer to drive the rollout turn by turn: the *policy model being trained* generates each turn, and TRL captures its tokens and log-probs and applies the environment mask — exactly what `environment_factory` provides over a black-box `rollout_func`. An installed agent is opaque to this: it runs inside the container with its *own* model and only emits a trajectory after the fact, so there are no policy tokens or log-probs for the trainer to optimize, and the model under training is never invoked. A [`~trl.experimental.harbor.HarborEnv`] is therefore an external agent — its tool methods `exec` into the sandbox, but the loop, and the model under training, stay in TRL. | ||
|
|
||
| ## Selecting the base agent (harness) | ||
|
|
||
| The **base agent** is the harness — which tool methods the env exposes and how it submits. Select it with `agent=`: | ||
|
|
||
| ```python | ||
| HarborSpec(dataset, agent="bash") # built-in single-bash-tool harness | ||
| HarborSpec(dataset, agent="my_pkg.harnesses:JupyterEnv") # import path to your HarborEnv subclass | ||
| HarborSpec(dataset, agent="path/to/harness.py:JupyterEnv") # file path to your HarborEnv subclass | ||
| HarborSpec(dataset, agent=MyHarborEnv) # a HarborEnv subclass directly | ||
| ``` | ||
|
|
||
| The built-in `"bash"` harness ([`~trl.experimental.harbor.HarborBashEnv`]) exposes one `bash` tool and submits by writing `/workdir/answer.txt`. Two richer harnesses ship as examples — each in its own folder with a README listing its tools — under [`examples/scripts/harbor/harnesses/`](https://github.com/huggingface/trl/tree/main/examples/scripts/harbor/harnesses): | ||
|
|
||
| - [`jupyter/`](https://github.com/huggingface/trl/tree/main/examples/scripts/harbor/harnesses/jupyter) (`JupyterEnv`) — a stateful Python kernel (variables persist across cells) + a shell tool. | ||
| - [`terminal_notes/`](https://github.com/huggingface/trl/tree/main/examples/scripts/harbor/harnesses/terminal_notes) (`TerminalNotesEnv`) — 6 shell tools (incl. background processes) + a 4-tool persistent note toolkit. | ||
|
|
||
| ```python | ||
| HarborSpec(dataset, agent="examples/scripts/harbor/harnesses/jupyter/env.py:JupyterEnv") | ||
| ``` | ||
|
|
||
| To write your own harness, subclass [`~trl.experimental.harbor.HarborEnv`] and add tool methods — every public method becomes a tool (TRL discovers them with `inspect.getmembers`), so give each a typed signature and a docstring (used to build the tool schema). Keep helpers underscore-prefixed. Use `self._exec(cmd)` to run shell commands in the sandbox, and set `PROMPT_SUFFIX` to append harness guidance to the task instruction: | ||
|
|
||
| ```python | ||
| from trl.experimental.harbor import HarborEnv | ||
|
|
||
| class GrepEnv(HarborEnv): | ||
| PROMPT_SUFFIX = "\n\nUse `grep` and `read_file`. Submit by writing /workdir/answer.txt." | ||
|
|
||
| def grep(self, pattern: str, path: str) -> str: | ||
| """Search for `pattern` under `path`. | ||
|
|
||
| Args: | ||
| pattern: The regex to search for. | ||
| path: The file or directory to search. | ||
| """ | ||
| return self._exec(f"grep -rn {pattern!r} {path!r}") | ||
| ``` | ||
|
|
||
| ## The sandbox backend | ||
|
|
||
| `environment_type` is passed straight through to Harbor (not validated by TRL): | ||
|
|
||
| ```python | ||
| HarborSpec(dataset, environment_type="e2b") # cloud sandbox (offloads provisioning), needs E2B_API_KEY | ||
| HarborSpec(dataset, environment_type="docker") # default; needs a local Docker daemon | ||
| ``` | ||
|
|
||
| `e2b` is recommended for cluster training: only `environment.exec` crosses into the cloud sandbox, so the GPUs stay dedicated to the policy and you can run many rollouts concurrently. | ||
|
|
||
| ## Reward functions | ||
|
|
||
| `spec.reward_funcs` defaults to an outcome reward — per rollout it reads the Harbor verifier's scalar (`env.reward`), computed once after the rollout by running the task's `tests/` verifier in the sandbox. For a custom reward, write a regular TRL reward function: | ||
|
|
||
| ```python | ||
| def my_reward(environments, **kwargs) -> list[float]: | ||
| return [env.reward for env in environments] | ||
| ``` | ||
|
|
||
| ## HarborSpec | ||
|
|
||
| [[autodoc]] trl.experimental.harbor.HarborSpec | ||
|
|
||
| ## Limitations | ||
|
|
||
| - The integration is in `trl.experimental` — APIs may change. Set `TRL_EXPERIMENTAL_SILENCE=1` to silence the warning in CI logs. | ||
| - Harbor's async sandbox client is bound to one event loop, so each env drives start/exec/verify synchronously on its own loop; sandbox provisioning is therefore sequential across the generation batch (cloud backends like `e2b` mitigate the per-sandbox cost). | ||
| - A single `HarborSpec` covers one task suite + one harness; multi-suite training is not supported yet. | ||
|
|
||
| ## Reference | ||
|
|
||
| - [Harbor framework](https://www.harborframework.com) | ||
| - [Harbor RL training docs](https://www.harborframework.com/docs/training-workflows/rl) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,131 @@ | ||
| # Copyright 2020-2026 The HuggingFace Team. All rights reserved. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| # /// script | ||
| # dependencies = ["trl[vllm,harbor]", "harbor[e2b]"] # harbor[e2b]: the --env e2b backend (swap per --env) | ||
| # /// | ||
|
|
||
| """GRPO training against a Harbor task suite. | ||
|
|
||
| One ``HarborSpec`` fans out into TRL's three slots — ``.train_dataset`` / ``.environment_factory`` / | ||
| ``.reward_funcs``. The base agent (harness) is selected with ``--agent``: the built-in ``"bash"``, or | ||
| an import/file path to a custom ``HarborEnv`` subclass. Harbor runs in-process, so install with | ||
| ``trl[harbor]`` (Python >= 3.12) plus the chosen sandbox backend's extra (``pip install "harbor[e2b]"`` | ||
| for the recommended ``--env e2b``; the ``docker`` backend just needs a reachable Docker daemon). | ||
|
|
||
| Usage (server vLLM, single-node 2+2 GPU split): | ||
|
|
||
| ```sh | ||
| # Terminal 1 — vLLM | ||
| CUDA_VISIBLE_DEVICES=2,3 trl vllm-serve --model Qwen/Qwen3-4B --tensor-parallel-size 2 --port 8000 | ||
|
|
||
| # Terminal 2 — training | ||
| CUDA_VISIBLE_DEVICES=0,1 accelerate launch \ | ||
| --config_file examples/accelerate_configs/deepspeed_zero2.yaml --num_processes 2 \ | ||
| examples/scripts/harbor/data_agent.py \ | ||
| --vllm-mode server --vllm-server-base-url http://localhost:8000 --env e2b | ||
| ``` | ||
| """ | ||
|
|
||
| import argparse | ||
|
|
||
| from trl import GRPOConfig, GRPOTrainer | ||
| from trl.experimental.harbor import HarborSpec | ||
|
|
||
|
|
||
| def parse_args() -> argparse.Namespace: | ||
| parser = argparse.ArgumentParser(description="GRPO training against a Harbor task suite.") | ||
|
|
||
| parser.add_argument("--model", type=str, default="Qwen/Qwen3-4B") | ||
| parser.add_argument( | ||
| "--dataset", | ||
| type=str, | ||
| default="AdithyaSK/data_agent_rl_environment_train", | ||
| help="HF dataset repo id with a Harbor task tree, or a local path containing a tasks/ subtree.", | ||
| ) | ||
| parser.add_argument( | ||
| "--agent", | ||
| type=str, | ||
| default="bash", | ||
| help="Base agent/harness: built-in name ('bash'), import path 'pkg.mod:Class', or file 'path.py:Class'.", | ||
| ) | ||
| parser.add_argument( | ||
| "--env", | ||
| dest="environment_type", | ||
| type=str, | ||
| default="docker", | ||
| help="Harbor sandbox backend (docker, e2b, daytona, gke, modal, runloop, ...).", | ||
| ) | ||
| parser.add_argument("--num-tasks", type=int, default=64) | ||
|
|
||
| parser.add_argument("--learning-rate", type=float, default=1e-6) | ||
| parser.add_argument("--per-device-train-batch-size", type=int, default=8) | ||
| parser.add_argument("--num-generations", type=int, default=8) | ||
| parser.add_argument("--max-completion-length", type=int, default=2048) | ||
| parser.add_argument("--max-steps", type=int, default=50) | ||
| parser.add_argument("--max-tool-calling-iterations", type=int, default=25) | ||
|
|
||
| parser.add_argument("--vllm-mode", choices=("colocate", "server"), default="colocate") | ||
| parser.add_argument("--vllm-server-base-url", type=str, default="http://localhost:8000") | ||
| parser.add_argument("--vllm-gpu-memory-utilization", type=float, default=0.3) | ||
|
|
||
| parser.add_argument("--output-dir", type=str, default=None) | ||
| parser.add_argument("--report-to", type=str, default="none") | ||
|
|
||
| return parser.parse_args() | ||
|
|
||
|
|
||
| def main() -> None: | ||
| args = parse_args() | ||
|
|
||
| # One spec object — fans out into TRL's three slots. | ||
| spec = HarborSpec( | ||
| args.dataset, | ||
| agent=args.agent, | ||
| environment_type=args.environment_type, | ||
| num_tasks=args.num_tasks, | ||
| ) | ||
|
|
||
| config_kwargs: dict = dict( | ||
| learning_rate=args.learning_rate, | ||
| per_device_train_batch_size=args.per_device_train_batch_size, | ||
| num_generations=args.num_generations, | ||
| max_completion_length=args.max_completion_length, | ||
| max_steps=args.max_steps, | ||
| max_tool_calling_iterations=args.max_tool_calling_iterations, | ||
| chat_template_kwargs={"enable_thinking": False}, | ||
| log_completions=True, | ||
| use_vllm=True, | ||
| vllm_mode=args.vllm_mode, | ||
| report_to=[s.strip() for s in args.report_to.split(",") if s.strip() and s.strip() != "none"] or "none", | ||
| ) | ||
| if args.output_dir: | ||
| config_kwargs["output_dir"] = args.output_dir | ||
| if args.vllm_mode == "colocate": | ||
| config_kwargs["vllm_gpu_memory_utilization"] = args.vllm_gpu_memory_utilization | ||
| else: | ||
| config_kwargs["vllm_server_base_url"] = args.vllm_server_base_url | ||
|
|
||
| trainer = GRPOTrainer( | ||
| model=args.model, | ||
| args=GRPOConfig(**config_kwargs), | ||
| train_dataset=spec.train_dataset, | ||
| environment_factory=spec.environment_factory, | ||
| reward_funcs=spec.reward_funcs, | ||
| ) | ||
| trainer.train() | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| main() |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| # Copyright 2020-2026 The HuggingFace Team. All rights reserved. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from .bash import BashEnv | ||
| from .jupyter import JupyterEnv | ||
| from .terminal_notes import TerminalNotesEnv | ||
|
|
||
|
|
||
| __all__ = ["BashEnv", "JupyterEnv", "TerminalNotesEnv"] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| # `bash` harness | ||
|
|
||
| The minimal harness: a single shell tool. This is the **built-in** [`HarborBashEnv`](../../../../../trl/experimental/harbor/_env.py) (`trl.experimental.harbor.HarborBashEnv`); this folder just documents it and re-exports it as `BashEnv`. | ||
|
|
||
| ## Tools | ||
|
|
||
| | Tool | Signature | What it does | | ||
| |---|---|---| | ||
| | `bash` | `bash(command: str) -> str` | Run a shell command in the sandbox; returns combined stdout+stderr (truncated to 8k). Non-stateful between calls. | | ||
|
|
||
| ## Submission | ||
|
|
||
| No submit tool — write the answer to `/workdir/answer.txt`, e.g. `echo -n "<value>" > /workdir/answer.txt`. The task's verifier reads that file. | ||
|
|
||
| ## Use it | ||
|
|
||
| ```python | ||
| from trl.experimental.harbor import HarborSpec | ||
| spec = HarborSpec(dataset, agent="bash") # built-in name | ||
| ``` |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| # Copyright 2020-2026 The HuggingFace Team. All rights reserved. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from trl.experimental.harbor import HarborBashEnv as BashEnv # the built-in bash harness | ||
|
|
||
|
|
||
| __all__ = ["BashEnv"] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this should be ordered alphabetically
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done — reordered to harbor → openenv → openreward.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
sorry for the misunderstanding 😅 the whole Integrations section should be ordered alphabetically, not only the env integrations