visualwebarena massage (#248)

ServiceNow · Nov 7, 2024 · 8712f8d · 8712f8d
1 parent d53fd5e
commit 8712f8d
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/browsergym/experiments/src/browsergym/experiments/benchmark/utils.py b/browsergym/experiments/src/browsergym/experiments/benchmark/utils.py
@@ -1,10 +1,14 @@
+import logging
 import os
 from typing import Literal
 
+import gymnasium as gym
 import numpy as np
 
 from browsergym.experiments.loop import SEED_MAX, EnvArgs
 
+logger = logging.getLogger(__name__)
+
 
 def make_env_args_list_from_workarena_curriculum(
     level: Literal["l1", "l2", "l3"],
@@ -122,6 +126,32 @@ def prepare_backend(backend: str):
             default_instance = VisualWebArenaInstance()
             default_instance.full_reset()
 
+            vwa_massage_task_ids = [
+                0,  # classifieds
+                33,  # classifieds
+                150,  # classifieds
+                253,  # reddit
+                325,  # reddit
+                390,  # reddit
+                444,  # shopping
+                555,  # shopping
+                666,  # shopping
+            ]
+            for task_id in vwa_massage_task_ids:
+                gym_id = f"browsergym/visualwebarena.{task_id}"
+                logger.info(
+                    f"VisualWebArena instance massaging {task_id} / {len(vwa_massage_task_ids)} ({gym_id} reset)"
+                )
+                env = gym.make(gym_id)
+                try:
+                    env.reset()  # task setup and logging
+                except Exception as e:
+                    logger.warning(
+                        f"Error during VisualWebArena instance massaging ({gym_id} reset): {e}"
+                    )
+                finally:
+                    env.close()
+
         case "workarena":
             # register environments
             import browsergym.workarena