Turing/ Playwright-based CUA environment with multi-provider adapter support

[raveedturing]

Turing/ Browser Gym — Playwright-based CUA environment for web navigation training

What does this PR do?

Adds a complete Computer-Use Agent (CUA) training environment to NeMo-Gym. Models interact with live web pages via Playwright — receiving screenshots and executing browser actions (click, type, scroll, keypress, drag, navigate). The environment supports three CUA providers: OpenAI (computer-use-preview), Anthropic (claude-sonnet-4-20250514 / claude-opus-4-6), and Google Gemini (gemini-2.5-computer-use-preview-10-2025).

This PR adds:

• Resources server (resources_servers/browser_gym/) — browser session lifecycle, Playwright action execution, session pooling with async locking, graceful error propagation, session reaper for stale sessions, standalone prepare_data.py CLI for offline task fetching
• Browser agent (responses_api_agents/browser_agent/) — CUA loop orchestration, provider-agnostic adapter architecture (OpenAICUAAdapter + GenericCUAAdapter), token ID tracking for RL training, debug trajectory output
• Model servers (responses_api_models/turing_browser_agent_anthropic_model/, responses_api_models/turing_browser_agent_gemini_model/) — stateless translators that accept OpenAI Responses API format, translate to/from provider-native APIs, and return unified NeMoGymResponse
• Core library changes — rollout cache memory optimization (retain_result_strs), CUA schema types (NeMoGymAction union + SDK-subclassing computer_call / computer_call_output), OpenAI organization header support

---

Architecture

Package layout

resources_servers/browser_gym/
├── app.py                  # FastAPI server: seed_session, step, verify, close, dump_local_storage
├── browser_pool.py         # Playwright browser pool with async session management
├── schemas.py              # CUA schemas (BrowserAction, CUAStepResponse, CUATrajectory, etc.)
├── setup_playwright.py     # Auto-install Chromium on startup (skips if already present)
├── prepare_data.py         # CLI for fetching tasks from gym URL → JSONL
├── configs/browser_gym.yaml
├── data/example.jsonl      # 5 example tasks
├── data/example_rollouts.jsonl
├── data/.gitignore
├── tests/test_app.py       # 68 tests
└── tests/test_prepare_data.py  # 17 tests

responses_api_agents/browser_agent/
├── app.py                  # CUA loop: seed → model → step → verify
├── trajectory_writer.py    # Debug trajectory output (screenshots + conversation JSONL)
├── adapters/
│   ├── base.py             # BaseCUAAdapter ABC + token ID extraction
│   ├── openai_adapter.py   # OpenAI computer-use-preview adapter (server-side context)
│   └── generic_adapter.py  # Generic CUA adapter for Anthropic + Gemini (client-side context)
└── tests/test_app.py       # 100 tests (209 with parameterization)

responses_api_models/turing_browser_agent_anthropic_model/
├── app.py                  # Stateless translator: OpenAI format ↔ Anthropic Messages API
├── configs/turing_browser_agent_anthropic_model.yaml
└── tests/test_app.py       # 41 tests

responses_api_models/turing_browser_agent_gemini_model/
├── app.py                  # Stateless translator: OpenAI format ↔ Gemini generate_content API
├── configs/turing_browser_agent_gemini_model.yaml
└── tests/test_app.py       # 52 tests

Data flow

gym eval run (--input <jsonl>)
  → agent /run
    → resources_server /seed_session (start browser, navigate to start_url, return screenshot)
    → loop:
        → model server /v1/responses (send screenshot + context in OpenAI format, get actions)
        → resources_server /step (execute action, return new screenshot + error if any)
        → adapter maps OpenAI-format response → unified BrowserAction
    → resources_server /dump_local_storage (capture browser state)
    → resources_server /close (tear down browser session)
    → resources_server /verify (send localStorage to gym API, compute reward)
  → JSONL output (reward + full trajectory with screenshots)

Adapter pattern

All providers follow the same interface. The adapter manages context internally, and all API calls use the standardized OpenAI Responses API format — model servers translate to/from provider-native APIs:

class BaseCUAAdapter(ABC):
    async def initialize(task_prompt, screenshot_b64) -> CUAAdapterResponse
    async def step(screenshot_b64, action_result, action_error) -> CUAAdapterResponse
    def reset()

Provider	Adapter	Context Management	Model Server
OpenAI	OpenAICUAAdapter	Server-side via previous_response_id	openai_model (proxies to api.openai.com)
Anthropic (Sonnet/Opus)	GenericCUAAdapter	Client-side: turn-based trimming, tool pair validation, screenshot GC	turing_browser_agent_anthropic_model (stateless translator: OpenAI format ↔ Anthropic Messages API)
Gemini	GenericCUAAdapter	Client-side: paired-turn trimming, function pair validation, screenshot GC	turing_browser_agent_gemini_model (stateless translator: OpenAI format ↔ Gemini generate_content API)

On model-server naming (addresses the earlier review note): the two turing_browser_agent_* model servers are deliberately generic stateless translators — they accept OpenAI Responses API format, translate to/from the provider-native API, and return NeMoGymResponse, with no browser-gym-specific logic (see their module docstrings). They are reusable by any environment, not coupled to Browser Gym. We are intentionally keeping the turing_browser_agent_ directory prefix — it reflects where the work originated, not a functional dependency.

Browser pool

• Configurable pool of Chromium processes (browser_pool_size) with semaphore-bounded concurrency (max_concurrent_browsers)
• Async locking for session management (add/remove)
• Session reaper for stale sessions (configurable TTL)
• Rollback on session creation failure (releases browser resources on error)
• Per-action timeouts with graceful error propagation back to the model
• Auto-installs Chromium on server startup (skips if already present)

Error handling

Action errors (e.g., invalid key names in keypress) are caught at the Playwright layer in browser_pool.py, returned through CUAStepResponse.error, and forwarded to the model via adapter-specific tool result formatting:

• OpenAI (OpenAICUAAdapter): Error set in computer_call_output.output.current_url
• Anthropic / Gemini (GenericCUAAdapter): Error set in computer_call_output.output.current_url (same OpenAI format — model server handles native translation)

This allows models to self-correct without crashing the trajectory. If the adapter's initialize() call fails (e.g., API unreachable), an empty trajectory is returned gracefully instead of a 500 that would crash rollout collection.

Verification

Verification uses localStorage assertions. After the CUA agent completes a task:

1. Agent dumps the browser's localStorage
2. Resource server sends form data (taskId, localStorageDump, initialState, modelResponse) to the gym's /api/v1/get_actual_state endpoint (with retry logic for transient 502/503/504 errors)
3. Gym returns assertions — reward is 1.0 if all pass, 0.0 otherwise

---

Core framework changes

nemo_gym/rollout_collection.py — Rollout cache memory optimization

• Optional result_strs retention in _load_from_cache(): Preserves the historical four-value return contract (input_rows, rows, results, result_strs) but adds a keyword-only retain_result_strs: bool = True parameter. When the caller will not upload to WandB, it skips keeping a second serialized copy (result_strs) of every cached rollout and parses results directly — avoiding unnecessary memory for large base64-heavy trajectories.
• Gated at the call site: run_from_config() passes retain_result_strs=config.upload_rollouts_to_wandb, so the extra copy is only retained when it will actually be used. The WandB result_strs are rebuilt at the end from results (upstream's run_from_config WandB flow, unchanged by this PR).

nemo_gym/openai_utils.py — CUA schemas + organization header

• NeMoGymAction union: Added a cross-provider computer-use action union (~20 types) covering the SDK actions (click, double_click, drag, keypress, move, screenshot, scroll, type, wait) plus the normalized browser actions the Anthropic/Gemini adapters emit (goto, go_back, go_forward, new_tab, switch_tab, close_tab, open_web_browser, search, click_at, hover_at, scroll_at, scroll_document, type_text_at, zoom, hold_key, …).
• CUA schema types: Added NeMoGymResponseComputerToolCall and NeMoGymComputerCallOutput to the NeMoGymResponseInputItem union. They inherit the upstream SDK types (ResponseComputerToolCall / ComputerCallOutput) — so issubclass holds and upstream's schema-drift tests keep passing — and only widen the fields CUA needs: action → the NeMoGymAction union (a superset of the SDK action union), and output → a plain Dict[str, Any] for the browser agent's screenshot/result payload. pending_safety_checks keeps a default so callers may omit it.
• Organization header: Added optional organization field to NeMoGymAsyncOpenAI — sends Openai-Organization header when set.

responses_api_models/openai_model/app.py — Organization support

• Organization support: Added openai_organization config field, passed through to NeMoGymAsyncOpenAI.

responses_api_models/turing_browser_agent_gemini_model/requirements.txt — Missing transitive dependencies

• Added cryptography and cffi as explicit dependencies. These are required by google-auth (a transitive dependency of google-genai) but were not being resolved automatically by uv during venv creation, causing ModuleNotFoundError at server startup.

tests/unit_tests/test_rollout_collection.py — Comprehensive coverage for rollout collection changes

• Added a _load_from_cache assertion for the retain_result_strs=False path (memory-optimized load returns the same rows with an empty result_strs), alongside the existing four-tuple (input_rows, rows, results, result_strs) coverage.
• Added run_from_config edge case tests: resume_from_cache with missing files, num_samples_in_parallel semaphore creation.
• Added test for _call_aggregate_metrics skipping rows without agent_ref.

---

Configuration

env.yaml (required)

max_concurrent_browsers: 16
browser_pool_size: 4
cua_debug_trajectories: false

# At least one provider:
cua_openai_api_key: <YOUR_OPENAI_API_KEY>
cua_openai_org: <YOUR_OPENAI_ORG>
cua_anthropic_api_key: <YOUR_ANTHROPIC_API_KEY>
cua_gemini_api_key: <YOUR_GEMINI_API_KEY>

Running

# Start all servers (resource + 4 model + 4 agent = 9 servers)
gym env start --config resources_servers/browser_gym/configs/browser_gym.yaml

# Prepare data from gym URL (offline, before rollout collection)
python resources_servers/browser_gym/prepare_data.py \
  --gym-url https://your-gym-url.com \
  --output resources_servers/browser_gym/data/tasks.jsonl

# Collect rollouts from a JSONL file
gym eval run --no-serve \
  --agent browser_openai_agent \
  --input resources_servers/browser_gym/data/tasks.jsonl \
  --output results/cua_rollouts.jsonl \
  --num-repeats 1 \
  --max-output-tokens 16384 \
  --temperature 1.0

# Profile results (pass@1 / pass@k, per-category metrics)
gym eval profile \
  --inputs results/cua_rollouts_materialized_inputs.jsonl \
  --rollouts results/cua_rollouts.jsonl

Available agents: browser_openai_agent, browser_anthropic_sonnet_agent, browser_anthropic_opus_agent, browser_gemini_agent.

---

Testing

391 test functions across the new Browser Gym components, all passing, plus targeted additions to test_rollout_collection.py for the core-library changes:

Component	Tests	Covers
resources_servers/browser_gym	71 + 17	BrowserPool, BrowserAction, key normalization, verify endpoint, session reaper, shutdown, rollback, prepare_data CLI
responses_api_agents/browser_agent	210	OpenAICUAAdapter + GenericCUAAdapter parsers, CUA loop orchestration (happy path, init failure, step failure, browser crash, error propagation, multi-step, usage accumulation), AdapterFactory for all providers, token ID extraction, trajectory writer
responses_api_models/turing_browser_agent_gemini_model	52	Config, OpenAI ↔ Gemini translation, content serialization/deserialization, tool derivation, response mapping
responses_api_models/turing_browser_agent_anthropic_model	41	Config, OpenAI ↔ Anthropic translation, tool/beta derivation, message mapping, response translation, output_config
tests/unit_tests/test_rollout_collection	added	_load_from_cache four-tuple return + retain_result_strs=False memory-opt path, resume_from_cache edge cases, num_samples_in_parallel semaphore, aggregate-metrics skip logic (alongside the existing upstream suite)

Pre-commit hooks (ruff, ruff-format, verified flag, README table) all pass.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[bxyu-nvidia]

can we revert this backoff logic pls? if we are hitting rate limits, my suggestion is to lower the concurrency

[raveedturing]

Done.

[bxyu-nvidia]

what's the reason we need input_gym_url ?

[raveedturing]

Removed this, and introduced the prepare_data script for the gym.

[bxyu-nvidia]

if the issue is the size, let's fork on if W&B is available and only append to result_strs when it is

[raveedturing]

Done.

[bxyu-nvidia]

same as above, if we are being rate limited let's lower the concurrency via ++num_samples_in_parallel

[raveedturing]

Done

[bxyu-nvidia]

we need to strictly adhere to OpenAI responses format for this route

[raveedturing]

Done

[bxyu-nvidia]

what's the highest concurrency we've tested this at?

[raveedturing]

Tested at 20 concurrent sessions with OpenAI, Anthropic, and Gemini. Higher concurrency is gated by provider API rate limits (per-org RPM/TPM) — CUA trajectories are particularly token-heavy since every turn includes a full screenshot as a base64 image, and each task requires its own Chromium instance for the duration of the trajectory (5–25 min).

[bxyu-nvidia]

we can decide to make this generic or specific. if we want to make it generic, we need some major refactor on this implementation. if specific, let's rename this folder to something like "browser_gym_anthropic_model" or something

[raveedturing]

Went with the generic approach — both model servers are now stateless translators (OpenAI format ↔ provider-native API) with no browser-gym-specific logic, reusable for any future agent. On the agent side, Anthropic and Gemini both use a single GenericCUAAdapter that speaks pure OpenAI format — provider SDKs are only dependencies of the respective model servers.

[raveedturing]

In addition to the generic approach, the model server names have also been made specific, starting with turing-.

[bxyu-nvidia]

can replace these try/excepts by setting ++global_aiohttp_client_request_debug=true

[raveedturing]

Done — removed the try/except blocks and the _sanitize_error() utility.