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

test_agent.py

@@ -0,0 +1,27 @@
+import gym
+from smarts.zoo import registry
+import sys
+
+test_agent = registry.make_agent("zoo.policies:competition_agent-v0")
+test_agent_spec = registry.make("zoo.policies:competition_agent-v0")
+
+shared_configs = dict(
+    action_space="TargetPose",
+    img_meters=64,
+    img_pixels=256,
+    headless=True,
+    sumo_headless=True,
+)
+
+test_env_path = "smarts.env:multi-scenario-v0"
+test_senario = "1_to_2lane_left_turn_c"
+test_env = gym.make(test_env_path, scenario=test_senario, **shared_configs)
+test_env = test_agent_spec.adapt_env(test_env)
+
+for _ in range(10):
+    observations = test_env.reset()
+    actions = test_agent.act(observations)
+    print(actions)
+    observations, rewards, dones, infos = test_env.step(actions)
+
+test_env.close()

zoo/policies/__init__.py

@@ -1,8 +1,16 @@
+import sys
+import os
+import importlib.util
+import shutil
+import subprocess
+import logging
+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 +108,126 @@ 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)
+    comp_env_path = str(
+        os.path.join(policy_path, "competition_env")
+    )  # folder contains all competition environment
+    sub_env_path = os.path.join(
+        comp_env_path, f"{Path(policy_path).name}"
+    )  # folder contains single competition environment
+    req_file = os.path.join(
+        policy_path, "requirements.txt"
+    )  # path of the requiremnet file
+
+    if Path(sub_env_path).exists():
+        shutil.rmtree(sub_env_path)
+    Path.mkdir(Path(sub_env_path), parents=True, exist_ok=True)
+
+    try:
+        subprocess.check_call(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "install",
+                "-t",
+                sub_env_path,
+                "-r",
+                req_file,
+            ]
+        )
+        sys.path.append(sub_env_path)
+    except:
+        logging.error(
+            f"Failed to install requirement for Competition Agent in folder {Path(policy_path).name}"
+        )
+        raise
+
+    # insert submission path
+    while policy_path in sys.path:
+        sys.path.remove(policy_path)
+
+    sys.path.insert(0, policy_path)
+
+    # import policy module
+    policy_file_path = str(os.path.join(policy_path, "policy.py"))
+    policy_spec = importlib.util.spec_from_file_location(
+        "competition_policy", policy_file_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)
+
+    policy = policy_module.Policy()
+    wrappers = policy_module.submitted_wrappers()
+
+    from .competition_agent import CompetitionAgent
+
+    def env_wrapper(env):
+        import gym
+
+        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,
+            "policy": policy,
+        },
+        adapt_env=env_wrapper,
+        agent_builder=CompetitionAgent,
+    )
+
+    # delete competition policy module and remove related path
+    while policy_path in sys.path:
+        sys.path.remove(policy_path)  # prevent duplicate path remain in sys.path
+
+    while sub_env_path in sys.path:
+        sys.path.remove(sub_env_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 sub_env_path in module_path
+            ):
+                sys.modules.pop(key)
+
+    del policy_module
+
+    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,40 @@
+import subprocess
+import sys
+import os
+import logging
+import importlib.util
+import shutil
+
+from pathlib import Path, PurePath
+from smarts.core.agent import Agent
+
+
+class CompetitionAgent(Agent):
+    def __init__(self, policy_path, policy):
+        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))
+
+        self._policy = policy
+
+    def act(self, obs):
+        return self._policy.act(obs)
+
+    def close(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)