update

Author: erikchwang

src/agentlab/agents/vl_agent/config.py

@@ -41,8 +41,8 @@
     "ui_agent": UIAgentArgs(
         main_vl_model_args=VL_MODEL_ARGS_DICT["gpt_4o"],
         auxiliary_vl_model_args=VL_MODEL_ARGS_DICT["llama_32_11b"],
-        action_set_args=HighLevelActionSetArgs(subsets=["coord"]),
         ui_prompt_args=VL_PROMPT_ARGS_DICT["ui_prompt"],
+        action_set_args=HighLevelActionSetArgs(subsets=["coord"]),
         max_retry=4,
     )
 }

src/agentlab/agents/vl_agent/vl_agent/base.py

@@ -1,9 +1,15 @@
 from abc import ABC, abstractmethod
+from browsergym.core.action.highlevel import HighLevelActionSet
 from browsergym.experiments.benchmark import Benchmark
 from dataclasses import dataclass
 
 
 class VLAgent(ABC):
+    @property
+    @abstractmethod
+    def action_set(self) -> HighLevelActionSet:
+        raise NotImplementedError
+
     @abstractmethod
     def get_action(self, obs: dict) -> tuple[str, dict]:
         raise NotImplementedError

src/agentlab/agents/vl_agent/vl_agent/ui_agent.py

@@ -1,11 +1,13 @@
 from agentlab.llm.llm_utils import ParseError, retry
 from agentlab.llm.tracking import cost_tracker_decorator
+from browsergym.core.action.highlevel import HighLevelActionSet
 from browsergym.experiments.agent import AgentInfo
 from browsergym.experiments.benchmark import Benchmark
 from browsergym.experiments.benchmark.base import HighLevelActionSetArgs
 from browsergym.utils.obs import overlay_som
 from copy import copy, deepcopy
 from dataclasses import asdict, dataclass
+from functools import cache
 from typing import Optional
 from .base import VLAgent, VLAgentArgs
 from ..vl_model.base import VLModelArgs
@@ -17,21 +19,26 @@ def __init__(
         self,
         main_vl_model_args: VLModelArgs,
         auxiliary_vl_model_args: Optional[VLModelArgs],
-        action_set_args: HighLevelActionSetArgs,
         ui_prompt_args: UIPromptArgs,
+        action_set_args: HighLevelActionSetArgs,
         max_retry: int,
     ):
         self.main_vl_model = main_vl_model_args.make_model()
         if auxiliary_vl_model_args is None:
             self.auxiliary_vl_model = None
         else:
             self.auxiliary_vl_model = auxiliary_vl_model_args.make_model()
-        self.action_set = action_set_args.make_action_set()
         self.ui_prompt_args = ui_prompt_args
+        self.action_set_args = action_set_args
         self.max_retry = max_retry
         self.thoughts = []
         self.actions = []
 
+    @property
+    @cache
+    def action_set(self) -> HighLevelActionSet:
+        return self.action_set_args.make_action_set()
+
     @cost_tracker_decorator
     def get_action(self, obs: dict) -> tuple[str, dict]:
         ui_prompt = self.ui_prompt_args.make_prompt(
@@ -94,25 +101,27 @@ def obs_preprocessor(self, obs: dict) -> dict:
 class UIAgentArgs(VLAgentArgs):
     main_vl_model_args: VLModelArgs
     auxiliary_vl_model_args: Optional[VLModelArgs]
-    action_set_args: HighLevelActionSetArgs
     ui_prompt_args: UIPromptArgs
+    action_set_args: HighLevelActionSetArgs
     max_retry: int
 
     @property
+    @cache
     def agent_name(self) -> str:
         if self.auxiliary_vl_model_args is None:
             return f"UIAgent-{self.main_vl_model_args.model_name}"
         else:
             return f"UIAgent-{self.main_vl_model_args.model_name}-{self.auxiliary_vl_model_args.model_name}"
 
     def make_agent(self) -> UIAgent:
-        return UIAgent(
+        self.ui_agent = UIAgent(
             main_vl_model_args=self.main_vl_model_args,
             auxiliary_vl_model_args=self.auxiliary_vl_model_args,
-            action_set_args=self.action_set_args,
             ui_prompt_args=self.ui_prompt_args,
+            action_set_args=self.action_set_args,
             max_retry=self.max_retry,
         )
+        return self.ui_agent
 
     def prepare(self):
         self.main_vl_model_args.prepare()

src/agentlab/agents/vl_agent/vl_model/llama_model.py

@@ -2,6 +2,7 @@
 from accelerate.utils.modeling import load_checkpoint_in_model
 from agentlab.llm.llm_utils import AIMessage, Discussion
 from dataclasses import dataclass
+from functools import cache
 from transformers import AutoProcessor, MllamaForConditionalGeneration
 from typing import Optional
 from .base import VLModel, VLModelArgs
@@ -84,6 +85,7 @@ class LlamaModelArgs(VLModelArgs):
     device: Optional[str]
 
     @property
+    @cache
     def model_name(self) -> str:
         return self.model_path.split("/")[-1].replace("-", "_").replace(".", "")
 
@@ -113,13 +115,14 @@ def make_model(self) -> LlamaModel:
         else:
             llama_model.model = llama_model.model.to(self.device)
         llama_model.model.eval()
-        return llama_model
+        self.llama_model = llama_model
+        return self.llama_model
 
     def prepare(self):
         pass
 
     def close(self):
-        pass
+        del self.llama_model.model
 
     def set_reproducibility_mode(self):
         self.reproducibility_config = {"do_sample": False}

src/agentlab/agents/vl_agent/vl_model/openrouter_api_model.py

@@ -1,5 +1,6 @@
 from agentlab.llm.llm_utils import AIMessage, Discussion
 from dataclasses import dataclass
+from functools import cache
 from openai import OpenAI, RateLimitError
 from .base import VLModel, VLModelArgs
 import backoff
@@ -46,16 +47,18 @@ class OpenRouterAPIModelArgs(VLModelArgs):
     reproducibility_config: dict
 
     @property
+    @cache
     def model_name(self) -> str:
         return self.model_id.split("/")[-1].replace("-", "_").replace(".", "")
 
     def make_model(self) -> OpenRouterAPIModel:
-        return OpenRouterAPIModel(
+        self.openrouter_api_model = OpenRouterAPIModel(
             base_url=self.base_url,
             model_id=self.model_id,
             max_tokens=self.max_tokens,
             reproducibility_config=self.reproducibility_config,
         )
+        return self.openrouter_api_model
 
     def prepare(self):
         pass