Fix BrowserGym server reset (#840)

* fix: harden browsergym server

* fix: keep BrowserGym task switches atomic

Ensure failed BrowserGym task recreation leaves the previous live environment and task state intact so callers can retry safely.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Test User <test@example.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 3 people

envs/browsergym_env/server/app.py

@@ -14,6 +14,9 @@
 viewport_width = int(os.environ.get("BROWSERGYM_VIEWPORT_WIDTH", "1280"))
 viewport_height = int(os.environ.get("BROWSERGYM_VIEWPORT_HEIGHT", "720"))
 timeout = float(os.environ.get("BROWSERGYM_TIMEOUT", "10000"))
+include_screenshot = (
+    os.environ.get("BROWSERGYM_INCLUDE_SCREENSHOT", "false").lower() == "true"
+)
 port = int(os.environ.get("BROWSERGYM_PORT", "8000"))
 
 max_concurrent = int(os.environ.get("MAX_CONCURRENT_ENVS", "8"))
@@ -27,6 +30,7 @@
         viewport_width=viewport_width,
         viewport_height=viewport_height,
         timeout=timeout,
+        include_screenshot=include_screenshot,
     ),
     BrowserGymAction,
     BrowserGymObservation,

envs/browsergym_env/server/browsergym_environment.py

@@ -24,6 +24,27 @@
 logger = logging.getLogger(__name__)
 
 
+def _to_jsonable(value: Any) -> Any:
+    """Convert BrowserGym payload values into JSON-safe Python objects."""
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    if isinstance(value, dict):
+        return {str(key): _to_jsonable(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_to_jsonable(item) for item in value]
+    if hasattr(value, "tolist"):
+        try:
+            return _to_jsonable(value.tolist())
+        except Exception:  # noqa: BLE001 - third-party array-likes
+            pass
+    if hasattr(value, "item"):
+        try:
+            return _to_jsonable(value.item())
+        except Exception:  # noqa: BLE001 - third-party scalar-likes
+            pass
+    return str(value)
+
+
 def _get_axtree_txt(obs: Dict[str, Any]) -> str:
     """Extract accessibility tree text from BrowserGym observation.
 
@@ -106,6 +127,7 @@ def __init__(
         viewport_width: int = 1280,
         viewport_height: int = 720,
         timeout: float = 10000.0,
+        include_screenshot: bool = False,
         **gym_kwargs: Any,
     ):
         """Initialize the BrowserGym environment.
@@ -128,14 +150,9 @@ def __init__(
         self.viewport_width = viewport_width
         self.viewport_height = viewport_height
         self.timeout = timeout
+        self.include_screenshot = include_screenshot
         self.gym_kwargs = dict(gym_kwargs)
 
-        # Build environment ID
-        if task_name:
-            self.env_id = f"browsergym/{benchmark}.{task_name}"
-        else:
-            self.env_id = f"browsergym/{benchmark}"
-
         # force import the benchmark module
         benchmark_modules = {
             "miniwob": "browsergym.miniwob",
@@ -158,23 +175,8 @@ def __init__(
             )
             raise ValueError(message) from import_error
 
-        # Create the BrowserGym environment
-        try:
-            self.gym_env = gym.make(
-                self.env_id,
-                headless=headless,
-                viewport={"width": viewport_width, "height": viewport_height},
-                timeout=timeout,
-                **self.gym_kwargs,
-            )
-        except Exception as e:  # noqa: BLE001 - gym.make
-            message = (
-                "Failed to create BrowserGym environment "
-                f"'{self.env_id}': {e}\n"
-                "Make sure the benchmark package is installed "
-                f"(e.g., pip install browsergym-{benchmark})."
-            )
-            raise ValueError(message) from e
+        self.env_id = self._build_env_id(task_name)
+        self.gym_env = self._create_gym_env(self.env_id)
 
         # State tracking
         self._state = BrowserGymState(
@@ -187,6 +189,51 @@ def __init__(
         self._last_obs: Optional[Dict[str, Any]] = None
         self._last_info: Optional[Dict[str, Any]] = None
 
+    def _build_env_id(self, task_name: Optional[str]) -> str:
+        if task_name:
+            return f"browsergym/{self.benchmark}.{task_name}"
+        return f"browsergym/{self.benchmark}"
+
+    def _create_gym_env(self, env_id: str):
+        try:
+            return gym.make(
+                env_id,
+                headless=self.headless,
+                viewport={
+                    "width": self.viewport_width,
+                    "height": self.viewport_height,
+                },
+                timeout=self.timeout,
+                **self.gym_kwargs,
+            )
+        except Exception as e:  # noqa: BLE001 - gym.make
+            message = (
+                "Failed to create BrowserGym environment "
+                f"'{env_id}': {e}\n"
+                "Make sure the benchmark package is installed "
+                f"(e.g., pip install browsergym-{self.benchmark})."
+            )
+            raise ValueError(message) from e
+
+    def _switch_task_if_needed(self, requested_task: str) -> None:
+        if requested_task == (self.task_name or ""):
+            return
+
+        next_task_name = requested_task or None
+        next_env_id = self._build_env_id(next_task_name)
+        next_gym_env = self._create_gym_env(next_env_id)
+
+        previous_gym_env = getattr(self, "gym_env", None)
+        self.task_name = next_task_name
+        self.env_id = next_env_id
+        self.gym_env = next_gym_env
+
+        if previous_gym_env is not None:
+            try:
+                previous_gym_env.close()
+            except Exception as exc:  # noqa: BLE001 - browsergym/playwright cleanup
+                logger.warning("BrowserGym cleanup failed during task switch: %s", exc)
+
     def reset(
         self,
         seed: Optional[int] = None,
@@ -202,11 +249,14 @@ def reset(
             Initial observation for the task
         """
         # Generate new episode ID
+        resolved_task_name = task_name or self.task_name or ""
+        self._switch_task_if_needed(resolved_task_name)
+
         self._state = BrowserGymState(
             episode_id=str(uuid4()),
             step_count=0,
             benchmark=self.benchmark,
-            task_name=task_name or self.task_name or "",
+            task_name=resolved_task_name,
         )
 
         # Reset options
@@ -343,7 +393,9 @@ def _create_observation(
         self._state.goal = goal
 
         # Extract additional observation modalities
-        screenshot = obs.get("screenshot") if isinstance(obs, dict) else None
+        screenshot = None
+        if self.include_screenshot and isinstance(obs, dict):
+            screenshot = _to_jsonable(obs.get("screenshot"))
 
         # Extract last_action_error from obs (BrowserGym includes this)
         last_action_error = False
@@ -357,11 +409,11 @@ def _create_observation(
         if isinstance(obs, dict):
             # Include useful fields but exclude large raw objects
             browsergym_metadata["browsergym_obs"] = {
-                k: v
+                k: _to_jsonable(v)
                 for k, v in obs.items()
                 if k not in ("dom_object", "axtree_object", "screenshot")
             }
-        browsergym_metadata["browsergym_info"] = info
+        browsergym_metadata["browsergym_info"] = _to_jsonable(info)
 
         return BrowserGymObservation(
             text=text,

tests/envs/test_browsergym_observation_serialization.py

@@ -0,0 +1,64 @@
+"""Unit tests for BrowserGym observation serialization."""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+pytest.importorskip("gymnasium", reason="gymnasium is not installed")
+
+from envs.browsergym_env.models import BrowserGymState
+from envs.browsergym_env.server.browsergym_environment import BrowserGymEnvironment
+
+
+class _ArrayLike:
+    def tolist(self):
+        return [[[255, 0, 0]]]
+
+
+class _ScalarLike:
+    def item(self):
+        return 3.5
+
+
+def test_create_observation_serializes_array_like_fields() -> None:
+    env = BrowserGymEnvironment.__new__(BrowserGymEnvironment)
+    env.include_screenshot = True
+    env._state = BrowserGymState()
+
+    result = BrowserGymEnvironment._create_observation(
+        env,
+        obs={
+            "goal": "Pick the highlighted option",
+            "url": "http://example.test",
+            "screenshot": _ArrayLike(),
+            "extra": {"score": _ScalarLike()},
+        },
+        info={"reward_estimate": _ScalarLike()},
+        done=False,
+        reward=0.0,
+    )
+
+    assert result.screenshot == [[[255, 0, 0]]]
+    assert result.metadata["browsergym_obs"]["extra"]["score"] == 3.5
+    assert result.metadata["browsergym_info"]["reward_estimate"] == 3.5
+
+
+def test_create_observation_omits_screenshot_by_default() -> None:
+    env = BrowserGymEnvironment.__new__(BrowserGymEnvironment)
+    env.include_screenshot = False
+    env._state = BrowserGymState()
+
+    result = BrowserGymEnvironment._create_observation(
+        env,
+        obs={"goal": "Pick the highlighted option", "screenshot": _ArrayLike()},
+        info={},
+        done=False,
+        reward=0.0,
+    )
+
+    assert result.screenshot is None

tests/envs/test_browsergym_task_switch.py

@@ -0,0 +1,124 @@
+"""Unit tests for BrowserGym task switching behavior."""
+
+from __future__ import annotations
+
+import os
+import sys
+import types
+from typing import Any
+
+import pytest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+
+try:
+    import gymnasium  # noqa: F401
+except ModuleNotFoundError:
+    sys.modules["gymnasium"] = types.SimpleNamespace(
+        make=lambda *_args, **_kwargs: None
+    )
+    _INSERTED_GYMNASIUM_STUB = True
+else:
+    _INSERTED_GYMNASIUM_STUB = False
+
+from envs.browsergym_env.server import browsergym_environment
+from envs.browsergym_env.server.browsergym_environment import BrowserGymEnvironment
+
+if _INSERTED_GYMNASIUM_STUB:
+    sys.modules.pop("gymnasium", None)
+
+
+class _FakeGymEnv:
+    def __init__(self, env_id: str) -> None:
+        self.env_id = env_id
+        self.closed = False
+        self.reset_calls: list[dict[str, Any]] = []
+
+    def reset(self, **kwargs: Any):
+        self.reset_calls.append(dict(kwargs))
+        return {"goal": f"goal for {self.env_id}", "url": "http://example.test"}, {}
+
+    def close(self) -> None:
+        self.closed = True
+
+
+def test_reset_rebuilds_browsergym_env_when_task_changes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    created_envs: list[tuple[str, dict[str, Any], _FakeGymEnv]] = []
+
+    monkeypatch.setattr(
+        browsergym_environment.importlib,
+        "import_module",
+        lambda _name: object(),
+    )
+
+    def fake_make(env_id: str, **kwargs: Any) -> _FakeGymEnv:
+        env = _FakeGymEnv(env_id)
+        created_envs.append((env_id, dict(kwargs), env))
+        return env
+
+    monkeypatch.setattr(browsergym_environment.gym, "make", fake_make)
+
+    env = BrowserGymEnvironment(
+        benchmark="miniwob",
+        task_name="click-test",
+        headless=True,
+    )
+    original_env = env.gym_env
+
+    observation = env.reset(task_name="enter-text", seed=7)
+
+    assert [env_id for env_id, _, _ in created_envs] == [
+        "browsergym/miniwob.click-test",
+        "browsergym/miniwob.enter-text",
+    ]
+    assert original_env.closed is True
+    assert env.gym_env is created_envs[1][2]
+    assert env.task_name == "enter-text"
+    assert env.state.task_name == "enter-text"
+    assert created_envs[1][2].reset_calls == [{"seed": 7}]
+    assert observation.goal == "goal for browsergym/miniwob.enter-text"
+
+
+def test_failed_task_switch_keeps_previous_browsergym_env(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    fail_enter_text = True
+    created_envs: list[_FakeGymEnv] = []
+
+    monkeypatch.setattr(
+        browsergym_environment.importlib,
+        "import_module",
+        lambda _name: object(),
+    )
+
+    def fake_make(env_id: str, **_kwargs: Any) -> _FakeGymEnv:
+        if env_id == "browsergym/miniwob.enter-text" and fail_enter_text:
+            raise RuntimeError("task is unavailable")
+        env = _FakeGymEnv(env_id)
+        created_envs.append(env)
+        return env
+
+    monkeypatch.setattr(browsergym_environment.gym, "make", fake_make)
+
+    env = BrowserGymEnvironment(benchmark="miniwob", task_name="click-test")
+    original_env = env.gym_env
+
+    with pytest.raises(ValueError, match="browsergym/miniwob.enter-text"):
+        env.reset(task_name="enter-text")
+
+    assert env.gym_env is original_env
+    assert original_env.closed is False
+    assert env.task_name == "click-test"
+    assert env.env_id == "browsergym/miniwob.click-test"
+    assert env.state.task_name == "click-test"
+
+    fail_enter_text = False
+    observation = env.reset(task_name="enter-text")
+
+    assert original_env.closed is True
+    assert env.gym_env is created_envs[-1]
+    assert env.task_name == "enter-text"
+    assert env.env_id == "browsergym/miniwob.enter-text"
+    assert observation.goal == "goal for browsergym/miniwob.enter-text"