improve profiling and delays

Author: recursix

src/agentlab/experiments/loop.py

@@ -48,6 +48,7 @@ class EnvArgs(DataClassJsonMixin):
     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
+    pre_observation_delay: float = 0.5  # seconds, wait for JS events to be fired
 
     def make_env(
         self, action_mapping, exp_dir, exp_task_kwargs: dict = {}, use_raw_page_output=True
@@ -71,6 +72,8 @@ def make_env(
             extra_kwargs["viewport"] = self.viewport
         if self.slow_mo is not None:
             extra_kwargs["slow_mo"] = self.slow_mo
+        if self.pre_observation_delay is not None:
+            extra_kwargs["pre_observation_delay"] = self.pre_observation_delay
         if self.storage_state:
             extra_kwargs["pw_context_kwargs"] = {"storage_state": self.storage_state}
         if self.task_kwargs is not None:
@@ -142,6 +145,12 @@ class StepTimestamps:
     env_stop: float = 0
     agent_start: float = 0
     agent_stop: float = 0
+    wait_for_page_loading_start: float = 0
+    wait_for_page_loading_stop: float = 0
+    validation_start: float = 0
+    validation_stop: float = 0
+    get_observation_start: float = 0
+    get_observation_stop: float = 0
 
 
 @dataclass
@@ -199,6 +208,12 @@ def from_step(self, env: gym.Env, action: str, obs_preprocessor: callable):
         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"]
+        t.wait_for_page_loading_start = env_info.get("wait_for_page_loading_start", None)
+        t.wait_for_page_loading_stop = env_info.get("wait_for_page_loading_stop", None)
+        t.validation_start = env_info.get("validation_start", None)
+        t.validation_stop = env_info.get("validation_stop", None)
+        t.get_observation_start = env_info.get("get_observation_start", None)
+        t.get_observation_stop = env_info.get("get_observation_stop", None)
 
         if obs_preprocessor:
             self.obs = obs_preprocessor(self.obs)
@@ -447,6 +462,10 @@ def run(self):
                 logger.debug("Sending action to environment.")
                 step_info.from_step(env, action, obs_preprocessor=agent.obs_preprocessor)
                 logger.debug("Environment stepped.")
+                if step_info.is_done:
+                    logger.debug(
+                        f"Episode done: terminated: {step_info.terminated}, truncated: {step_info.truncated}."
+                    )
 
         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}"