Repository navigation
Add harness session runtime and BrowserGym adapter for training and evaluation #471
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
Closed
Closed
Changes from 16 commits
Commits
Show all changes
20 commits
Select commit
Hold shift + click to select a range
52902a4
Add harness session runtime for training and evaluation
burtenshaw 774f445
Handle MCP tool errors and lazy-load BrowserGym exports
burtenshaw 3dc09ce
Refine experimental harness APIs and reward forwarding
burtenshaw fc0188e
Merge branch 'main' into feature/harness-interface
burtenshaw f67070d
docs: mention environment_factory in harness tutorial
burtenshaw 896c8c0
test: cover harness review regressions
burtenshaw 8637512
chore: merge main
burtenshaw 5414298
feat(harness): rollout collection + openenv collect CLI (#560)
sergiopaniego 032b3d3
fix: address browsergym harness review comments
burtenshaw 03c9261
fix: format collect files
burtenshaw 1448944
envs/opencode_env: OpenCode coding-agent harness primitive (#603)
adithya-s-k 88240b7
Merge branch 'main' into feature/harness-interface
burtenshaw dc451de
fix: address ci docs and formatting
burtenshaw d4032a1
feat(tutorials): add SFT warm-up tutorial with reasoning_gym collect …
sergiopaniego 004168d
Merge branch 'main' into feature/harness-interface
burtenshaw b943eb1
Merge branch 'main' into feature/harness-interface
burtenshaw da5a107
fix: address harness review feedback
burtenshaw 99cffbb
fix: format harness lint
burtenshaw 69a6fd7
fix: address harness review comments
burtenshaw 627948c
Merge branch 'main' into feature/harness-interface
burtenshaw File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,2 @@ | ||
| ```{include} ../../../envs/opencode_env/README.md | ||
| ``` |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,129 @@ | ||
| # BrowserGym Harness Rollouts | ||
|
|
||
| This tutorial shows how to drive BrowserGym through the OpenEnv harness runtime | ||
| when a trainer needs to keep token sampling, logprobs, and reward assignment | ||
| inside the training loop. | ||
|
|
||
| ```{note} | ||
| Use this pattern for tool-driven BrowserGym rollouts. For a standard | ||
| `reset()` / `step()` GRPO flow, keep using the Wordle GRPO tutorial. | ||
| ``` | ||
|
|
||
| ## What You'll Build | ||
|
|
||
| - A BrowserGym session factory that creates one environment client per rollout. | ||
| - A harness rollout function that TRL can call during training. | ||
| - A model-step wrapper that converts generated BrowserGym action text into | ||
| structured tool calls. | ||
|
|
||
| ## Install Dependencies | ||
|
|
||
| Install OpenEnv, TRL, and the BrowserGym environment package: | ||
|
|
||
| ```bash | ||
| pip install -U "trl[vllm]" peft trackio kernels | ||
| pip install -U git+https://github.com/meta-pytorch/OpenEnv.git | ||
| pip install -U "openenv-browsergym @ git+https://huggingface.co/spaces/openenv/browsergym_env" | ||
| ``` | ||
|
|
||
| ## Build The Session Factory | ||
|
|
||
| `BrowserGymSessionFactory` adapts a BrowserGym client into the harness | ||
| `ResourceSession` interface. If your training setup already has an | ||
| `environment_factory`, pass that factory as `client_factory` so every rollout | ||
| gets a fresh environment instance. | ||
|
|
||
| ```python | ||
| from browsergym_env import BrowserGymEnv | ||
| from browsergym_env.harness import BrowserGymSessionFactory | ||
|
|
||
| space_url = "https://openenv-browsergym-env.hf.space" | ||
|
|
||
|
|
||
| def environment_factory(): | ||
| return BrowserGymEnv(base_url=space_url) | ||
|
|
||
|
|
||
| session_factory = BrowserGymSessionFactory( | ||
| client_factory=environment_factory, | ||
| default_task="click-test", | ||
| ) | ||
| ``` | ||
|
|
||
| The session exposes BrowserGym actions such as `click`, `fill`, `send_keys`, | ||
| `scroll`, and `noop` as MCP-style tools while still executing the corresponding | ||
| BrowserGym action strings under the hood. | ||
|
|
||
| ## Wrap TRL Generation | ||
|
|
||
| The harness calls a `model_step` function for each turn. The model step should | ||
| use the trainer-owned generation path, then return a `ModelStepResult` with the | ||
| completion text, token ids, logprobs, and exactly one BrowserGym tool call. | ||
|
|
||
| ```python | ||
| from browsergym_env.harness import build_browsergym_action_tool_call | ||
| from openenv.core.harness import ModelStepResult | ||
| from openenv.core.llm_client import LLMResponse | ||
| from trl.experimental.openenv import generate_rollout_completions | ||
|
|
||
|
|
||
| def build_trl_browsergym_model_step(trainer, tokenizer): | ||
| def model_step(messages, tools, sampling): | ||
| del tools, sampling | ||
| prompt_text = tokenizer.apply_chat_template( | ||
| messages, | ||
| add_generation_prompt=True, | ||
| tokenize=False, | ||
| ) | ||
| rollout_output = generate_rollout_completions(trainer, [prompt_text])[0] | ||
| completion_text = rollout_output.get("text") or tokenizer.decode( | ||
| rollout_output["completion_ids"], | ||
| skip_special_tokens=True, | ||
| ) | ||
| tool_call = build_browsergym_action_tool_call(completion_text) | ||
| return ModelStepResult( | ||
| response=LLMResponse(content=completion_text, tool_calls=[tool_call]), | ||
| prompt_ids=list(rollout_output["prompt_ids"]), | ||
| completion_ids=list(rollout_output["completion_ids"]), | ||
| logprobs=list(rollout_output["logprobs"]), | ||
| ) | ||
|
|
||
| return model_step | ||
| ``` | ||
|
|
||
| In practice, you should add a small parser around the completion text so common | ||
| outputs like `Action: click('13')` are normalized before calling | ||
| `build_browsergym_action_tool_call`. | ||
|
|
||
| ## Create The Rollout Function | ||
|
|
||
| Pass the session factory, white-box harness adapter, and model-step builder to | ||
| `build_harness_rollout_func`: | ||
|
|
||
| ```python | ||
| from openenv.core.harness import ( | ||
| HarnessRunLimits, | ||
| MCPHarnessAdapter, | ||
| build_harness_rollout_func, | ||
| ) | ||
|
|
||
| rollout_func = build_harness_rollout_func( | ||
| session_factory=session_factory, | ||
| harness_adapter=MCPHarnessAdapter(), | ||
| model_step_builder=lambda trainer, session: build_trl_browsergym_model_step( | ||
| trainer, | ||
| tokenizer, | ||
| ), | ||
| limits=HarnessRunLimits(max_turns=10), | ||
| ) | ||
| ``` | ||
|
|
||
| The returned function accepts TRL prompts and a trainer, runs one harness-backed | ||
| BrowserGym episode per prompt, and returns `prompt_ids`, `completion_ids`, | ||
| `logprobs`, `env_reward`, and `verify_metrics`. | ||
|
|
||
| ## Full Example | ||
|
|
||
| See [`examples/browsergym_harness.py`](https://github.com/meta-pytorch/OpenEnv/blob/main/examples/browsergym_harness.py) | ||
| for a complete TRL-oriented helper that includes action normalization and a | ||
| ready-to-use `build_browsergym_rollout_func`. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.
The rest of the tutorials now have a
Open in Colabbadge with its corresponding notebook. Could be added here or as a follow-up PR.