weblinx prepare_backend() fix

browsergym/experiments/src/browsergym/experiments/benchmark/base.py

@@ -1,5 +1,6 @@
 import fnmatch
 import logging
+import os
 import typing
 from dataclasses import dataclass, field
 from typing import Literal, Optional
@@ -86,6 +87,7 @@ def __post_init__(self):
 
     def prepare_backends(self):
         for backend in self.backends:
+            logger.info(f"Preparing {backend} backend...")
             match backend:
                 case "miniwob":
                     # register environments
@@ -134,9 +136,28 @@ def prepare_backends(self):
                     # register environments
                     import weblinx_browsergym
 
+                    # pre-download all weblinx files
+                    cache_dir = os.environ.get("BROWSERGYM_WEBLINX_CACHE_DIR", None)
+
+                    assert (
+                        cache_dir
+                    ), f"Environment variable BROWSERGYM_WEBLINX_CACHE_DIR is missing or empty, required to prepare the weblinx backend."
+
+                    all_tasks = []
+                    for split in ("train", "valid", "test_iid"):
+                        all_tasks.extend(
+                            weblinx_browsergym.list_tasks(split=split, cache_dir=cache_dir)
+                        )
+                    demo_ids = weblinx_browsergym.get_unique_demo_ids(tasks=all_tasks)
+                    weblinx_browsergym.download_and_unzip_demos(
+                        demo_ids=demo_ids, cache_dir=cache_dir
+                    )
+
                 case _:
                     raise ValueError(f"Unknown benchmark backend {repr(backend)}")
 
+            logger.info(f"{backend} backend ready")
+
     def subset_from_split(self, split: Literal["train", "valid", "test"]):
         split_column = "browsergym_split"