Competition agent integration with zoo

smarts/zoo/agent_spec.py

@@ -74,6 +74,8 @@ class AgentSpec:
     """An adaptor that allows shaping of the reward (default lambda obs, reward: reward)"""
     info_adapter: Callable = lambda obs, reward, info: info
     """An adaptor that allows shaping of info (default lambda obs, reward, info: info)"""
+    adapt_env: Optional[Callable]= lambda env: env
+    """And adaptor that allows configurating the env (default lambda env:env)"""
 
     def __post_init__(self):
         # make sure we can pickle ourselves

zoo/policies/__init__.py

@@ -1,8 +1,13 @@
+import sys
+import os
+import importlib.util
+from pathlib import Path
 from typing import Any, Dict
 
 from smarts.core.agent_interface import AgentInterface, AgentType
 from smarts.core.controllers import ActionSpaceType
 from smarts.zoo.agent_spec import AgentSpec
+from smarts.env.multi_scenario_env import resolve_agent_interface
 from smarts.zoo.registry import make, register
 
 from .keep_lane_agent import KeepLaneAgent
@@ -100,3 +105,88 @@ def human_keyboard_entrypoint(*arg, **kwargs):
 
 
 register(locator="human-in-the-loop-v0", entry_point=human_keyboard_entrypoint)
+
+
+def load_config(path):
+    import yaml
+
+    config = None
+    if path.exists():
+        with open(path, "r") as file:
+            config = yaml.safe_load(file)
+    return config
+
+
+root_path = str(Path(__file__).absolute().parents[2])
+
+
+def competition_entry(**kwargs):
+    policy_path = kwargs.get("policy_path", None)
+    env_path = os.path.join(root_path, f"competition_env/{Path(policy_path).name}")
+
+    from .competition_agent import CompetitionAgent
+
+    def env_wrapper(env):
+        import gym
+
+        # insert policy path
+        if policy_path in sys.path:
+            sys.path.remove(policy_path)
+
+        sys.path.insert(0, policy_path)
+
+        # import policy.py module
+        wrapper_path = str(os.path.join(policy_path, "policy.py"))
+        wrapper_spec = importlib.util.spec_from_file_location(
+            "competition_wrapper", wrapper_path
+        )
+        wrapper_module = importlib.util.module_from_spec(wrapper_spec)
+        sys.modules["competition_wrapper"] = wrapper_module
+        if wrapper_spec:
+            wrapper_spec.loader.exec_module(wrapper_module)
+
+        wrappers = wrapper_module.submitted_wrappers()
+
+        # delete competition wrapper module and remove path
+        while policy_path in sys.path:
+            sys.path.remove(policy_path)  # prevent duplicate path remain in sys.path
+
+        # remove all modules related to policy_path
+        for key, module in list(sys.modules.items()):
+            if "__file__" in dir(module):
+                module_path = module.__file__
+                if module_path and (
+                    policy_path in module_path or env_path in module_path
+                ):
+                    sys.modules.pop(key)
+        del wrapper_module
+
+        env = gym.Wrapper(env)
+        for wrapper in wrappers:
+            env = wrapper(env)
+
+        return env
+
+    config = load_config(Path(os.path.join(policy_path, "config.yaml")))
+
+    spec = AgentSpec(
+        interface=resolve_agent_interface(
+            img_meters=int(config["img_meters"]),
+            img_pixels=int(config["img_pixels"]),
+            action_space="TargetPose",
+        ),
+        agent_params={
+            "policy_path": policy_path,
+        },
+        adapt_env=env_wrapper,
+        agent_builder=CompetitionAgent,
+    )
+
+    return spec
+
+
+register(
+    "competition_agent-v0",
+    entry_point=competition_entry,
+    policy_path=os.path.join(root_path, "competition/track1/submission"),
+)

zoo/policies/competition_agent.py

@@ -0,0 +1,91 @@
+import subprocess
+import sys
+import os
+import importlib.util
+import shutil
+
+from pathlib import Path, PurePath
+from smarts.core.agent import Agent
+
+
+class CompetitionAgent(Agent):
+    def __init__(self, policy_path):
+        req_file = os.path.join(policy_path, "requirements.txt")
+
+        env_name = Path(policy_path).name  # name of the submission file
+        root_path = Path(__file__).parents[2]  # Smarts main path
+
+        self._policy_dir = policy_path
+        self._comp_env_path = str(os.path.join(root_path, "competition_env"))
+        self._sub_env_path = str(os.path.join(self._comp_env_path, env_name))
+
+        Path.mkdir(Path(self._comp_env_path), exist_ok=True)
+
+        if Path(self._sub_env_path).exists():
+            shutil.rmtree(self._sub_env_path)
+        Path.mkdir(Path(self._sub_env_path))
+
+        try:
+            subprocess.check_call(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "install",
+                    "-t",
+                    self._sub_env_path,
+                    "-r",
+                    req_file,
+                ]
+            )
+            sys.path.append(self._sub_env_path)
+        except:
+            sys.exit("Failed to install requirement for Competition Agent")
+
+        # insert submission path
+        sys.path.insert(0, policy_path)
+
+        # import policy module
+        self._policy_path = str(os.path.join(policy_path, "policy.py"))
+        policy_spec = importlib.util.spec_from_file_location(
+            "competition_policy", self._policy_path
+        )
+        policy_module = importlib.util.module_from_spec(policy_spec)
+        sys.modules["competition_policy"] = policy_module
+        if policy_spec:
+            policy_spec.loader.exec_module(policy_module)
+
+        self._policy = policy_module.Policy()
+
+        # delete competition policy module and remove path
+        while policy_path in sys.path:
+            sys.path.remove(policy_path)  # prevent duplication path remain in sys.path
+
+        # remove all modules related to the policy_path
+        for key, module in list(sys.modules.items()):
+            if "__file__" in dir(module):
+                module_path = module.__file__
+                if module_path and (
+                    self._policy_dir in module_path or self._sub_env_path in module_path
+                ):
+                    sys.modules.pop(key)
+        del policy_module
+
+    def act(self, obs):
+        return self._policy.act(obs)
+
+    def close_env(self, remove_all_env=False):
+        shutil.rmtree(str(self._sub_env_path))
+        while self._sub_env_path in sys.path:
+            sys.path.remove(self._sub_env_path)
+        while self._policy_dir in sys.path:
+            sys.path.remove(self._policy_dir)
+        for key, module in list(sys.modules.items()):
+            if "__file__" in dir(module):
+                module_path = module.__file__
+                if module_path and (
+                    self._policy_dir in module_path or self._sub_env_path in module_path
+                ):
+                    sys.modules.pop(key)
+        if remove_all_env:
+            shutil.rmtree(self._comp_env_path, ignore_errors=True)