selective action validation

Author: ollmer

browsergym/core/src/browsergym/core/env.py

@@ -427,7 +427,7 @@ def step(self, action: str) -> tuple[dict[str, Any], float, bool, bool, dict[str
         return self.post_step(info)
 
     def post_step(
-        self, info: dict[str, Any]
+        self, info: dict[str, Any], validate: bool = True
     ) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
         """
         Post step method, called after executing the action.
@@ -452,20 +452,27 @@ def post_step(
         # wait for the network to idle before extracting the observation, reward etc.
         self._wait_dom_loaded()
 
-        # after the action is executed, the active page might have changed
-        # perform a safety check
-        self._active_page_check()
-        logger.debug("Active page checked")
-
-        # if asked, wait for user message
-        self._wait_for_user_message()
-        logger.debug("User message done")
-
-        logger.debug("Initiating task validation")
-        # extract reward, done, user_message, info (task-specific)
-        reward, done, user_message, task_info = self._task_validate()
-        info["task_info"] = task_info
-        logger.debug("Task validation done")
+        if validate:
+            # after the action is executed, the active page might have changed
+            # perform a safety check
+            self._active_page_check()
+            logger.debug("Active page checked")
+
+            # if asked, wait for user message
+            self._wait_for_user_message()
+            logger.debug("User message done")
+
+            logger.debug("Initiating task validation")
+            # extract reward, done, user_message, info (task-specific)
+            reward, done, user_message, task_info = self._task_validate()
+            info["task_info"] = task_info
+            logger.debug("Task validation done")
+        else:
+            reward = 0
+            done = False
+            user_message = None
+            info["task_info"] = {}
+            logger.debug("Task validation skipped")
 
         # add any user message sent by the task to the chat
         if user_message:

browsergym/core/src/browsergym/utils/mcp_server.py

@@ -4,7 +4,7 @@
 import re
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Callable
 
 import gymnasium as gym
@@ -20,6 +20,7 @@ class BgymConfig:
     timeout_ms: int = 10000
     record_video_dir: str | None = None
     demo_mode: HighLevelActionSet.DemoMode = "default"
+    validate_actions: list[str] = field(default_factory=list)
 
 
 @dataclass
@@ -77,6 +78,13 @@ def get_cli_args():
         choices=ACTION_SUBSETS.keys(),
         help="Subset of actions to use",
     )
+    parser.add_argument(
+        "--validate_actions",
+        type=str,
+        nargs="+",
+        default=["click", "goto"],
+        help="Names of actions for which validation should be performed",
+    )
     args, _ = parser.parse_known_args()
     return args
 
@@ -88,6 +96,7 @@ def get_cli_args():
     timeout_ms=args.timeout_ms,
     record_video_dir=args.record_video_dir,
     demo_mode=args.demo_mode,
+    validate_actions=args.validate_actions,
 )
 
 
@@ -117,7 +126,7 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
 mcp = FastMCP("BrowserGym", lifespan=app_lifespan)
 
 
-def fn_wrapper(func: Callable):
+def fn_wrapper(func: Callable, validate: bool = True):
     async def decorator(*args, **kwargs):
         """
         Decorator to execute function from the action space in the context of the gym.
@@ -157,7 +166,7 @@ async def decorator(*args, **kwargs):
             if match:
                 info["action_exec_timeout"] = float(match.groups()[0]) / 1000
 
-        results = await asyncio.to_thread(gym.post_step, info)
+        results = await asyncio.to_thread(gym.post_step, info, validate)
         return results
 
     decorator.__wrapped__ = func  # type: ignore
@@ -167,7 +176,8 @@ async def decorator(*args, **kwargs):
 
 
 for fn in ACTION_SUBSETS[args.subset]:
-    mcp.add_tool(fn_wrapper(fn))
+    validate =fn.__name__ in config.validate_actions
+    mcp.add_tool(fn_wrapper(fn, validate))
 
 if __name__ == "__main__":
     mcp.run(transport="stdio")