Lazy importing of agents

deep_quoridor/src/agents/__init__.py

@@ -3,23 +3,10 @@
     "ActionLog",
     "Agent",
     "AgentRegistry",
-    "AlphaZeroOSAgent",
-    "DExpAgent",
-    "HumanAgent",
-    "MCTSAgent",
-    "NDexpAgent",
-    "RandomAgent",
     "ReplayAgent",
     "ReplayBuffer",
-    "SimpleAgent",
     "TrainableAgentParams",
 ]
-
-
-from agents.adapter_based_agents import Cnn3CAgent, CnnAgent, NDexpAgent
-from agents.alphazero import AlphaZeroAgent  # noqa: E402, F401
-from agents.alphazero_dexp import DAZAgent
-from agents.alphazero_os import AlphaZeroOSAgent  # noqa: E402, F401
 from agents.core import (  # noqa: E402, F401  # noqa: E402, F401
     AbstractTrainableAgent,
     ActionLog,
@@ -28,27 +15,20 @@
     ReplayBuffer,
     TrainableAgentParams,
 )
-from agents.dexp import DExpAgent  # noqa: E402
-from agents.greedy import GreedyAgent  # noqa: E402, F401
-from agents.human import HumanAgent  # noqa: E402, F401
-from agents.mcts import MCTSAgent  # noqa: E402, F401
-from agents.random import RandomAgent  # noqa: E402, F401
-from agents.replay import ReplayAgent  # noqa: E402, F401
-from agents.sb3_ppo import SB3PPOAgent  # noqa: E402, F401
-from agents.simple import SimpleAgent  # noqa: E402, F401
+from agents.replay import ReplayAgent
 
-AgentRegistry.register("alphazero", AlphaZeroAgent)
-AgentRegistry.register("alphazero_os", AlphaZeroOSAgent)
-AgentRegistry.register("cnn", CnnAgent)
-AgentRegistry.register("cnn3c", Cnn3CAgent)
-AgentRegistry.register("daz", DAZAgent)
-AgentRegistry.register("daz_mimic", DAZAgent.create_from_trained_instance)
-AgentRegistry.register("dexp", DExpAgent)
-AgentRegistry.register("dexp_mimic", DExpAgent.create_from_trained_instance)
-AgentRegistry.register("greedy", GreedyAgent)
-AgentRegistry.register("human", HumanAgent)
-AgentRegistry.register("mcts", MCTSAgent)
-AgentRegistry.register("ndexp", NDexpAgent)
-AgentRegistry.register("random", RandomAgent)
-AgentRegistry.register("simple", SimpleAgent)
-AgentRegistry.register("sb3ppo", SB3PPOAgent)
+AgentRegistry.register("alphazero", "AlphaZeroAgent", "agents.alphazero")
+AgentRegistry.register("alphazero_os", "AlphaZeroOSAgent", "agents.alphazero_os")
+AgentRegistry.register("cnn", "CnnAgent", "agents.adapter_based_agents")
+AgentRegistry.register("cnn3c", "Cnn3CAgent", "agents.adapter_based_agents")
+AgentRegistry.register("daz", "DAZAgent", "agents.alphazero_dexp")
+AgentRegistry.register("daz_mimic", "DAZAgent.create_from_trained_instance", "agents.alphazero_dexp")
+AgentRegistry.register("dexp", "DExpAgent", "agents.dexp")
+AgentRegistry.register("dexp_mimic", "DExpAgent.create_from_trained_instance", "agents.dexp")
+AgentRegistry.register("greedy", "GreedyAgent", "agents.greedy")
+AgentRegistry.register("human", "HumanAgent", "agents.human")
+AgentRegistry.register("mcts", "MCTSAgent", "agents.mcts")
+AgentRegistry.register("ndexp", "NDexpAgent", "agents.adapter_based_agents")
+AgentRegistry.register("random", "RandomAgent", "agents.random")
+AgentRegistry.register("simple", "SimpleAgent", "agents.simple")
+AgentRegistry.register("sb3ppo", "SB3PPOAgent", "agents.sb3_ppo")

deep_quoridor/src/agents/core/agent.py

@@ -1,8 +1,10 @@
+import importlib
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from typing import Optional, Type
 
 from quoridor import Action
-from utils import parse_subargs
+from utils import SubargsBase, parse_subargs
 
 
 class ActionLog:
@@ -113,9 +115,36 @@ def get_action(self, observation) -> int:
         pass
 
 
+@dataclass
+class AgentRegistryEntry:
+    class_name: str
+    module_name: str
+    agent_class: Optional[Type[Agent]] = None
+    params_class: Optional[Type[SubargsBase]] = None
+
+
 class AgentRegistry:
     agents = {}
 
+    @staticmethod
+    def get_registry_entry(agent_type: str) -> AgentRegistryEntry:
+        registry_entry = AgentRegistry.agents[agent_type]
+
+        if registry_entry.agent_class is None:
+            agent_module = importlib.import_module(registry_entry.module_name)
+            element_names = registry_entry.class_name.split(".")
+
+            # If the class_name is heirarchical, e.g. "foo.bar.Baz", then we need to gettatr
+            # the first element, then the second, etc. until we get to Baz.
+            registry_entry.agent_class = getattr(agent_module, element_names[0])
+            for element_name in element_names[1:]:
+                registry_entry.agent_class = getattr(registry_entry.agent_class, element_name)
+
+            if hasattr(registry_entry.agent_class, "params_class"):
+                registry_entry.params_class = registry_entry.agent_class.params_class()
+
+        return registry_entry
+
     @staticmethod
     def create(friendly_name: str, **kwargs) -> Agent:
         return AgentRegistry.agents[friendly_name](**kwargs)
@@ -126,19 +155,20 @@ def create_from_encoded_name(
     ) -> Agent:
         parts = encoded_name.split(":")
         agent_type = parts[0]
+        registry_entry = AgentRegistry.get_registry_entry(agent_type)
         if len(parts) == 2:
-            subargs_class = AgentRegistry.agents[agent_type].params_class()
-            if subargs_class is None:
+            if registry_entry.params_class is None:
                 raise ValueError(f"The agent {agent_type} doesn't support subarguments, but '{parts[1]}' was passed")
 
             if remove_training_args:
-                args_to_remove = subargs_class.training_only_params().difference(keep_args)
-                subargs = parse_subargs(parts[1], subargs_class, ignore_fields=args_to_remove)
+                args_to_remove = registry_entry.agent_class.training_only_params().difference(keep_args)
+                subargs = parse_subargs(parts[1], registry_entry.params_class, ignore_fields=args_to_remove)
             else:
-                subargs = parse_subargs(parts[1], subargs_class)
+                subargs = parse_subargs(parts[1], registry_entry.params_class)
+
             kwargs["params"] = subargs
 
-        return AgentRegistry.agents[agent_type](
+        return registry_entry.agent_class(
             board_size=env.board_size,
             max_walls=env.max_walls,
             observation_space=env.observation_space(None),
@@ -157,5 +187,5 @@ def names():
         return list(AgentRegistry.agents.keys())
 
     @staticmethod
-    def register(name: str, agent_class):
-        AgentRegistry.agents[name] = agent_class
+    def register(name: str, class_name: str, module_name: str):
+        AgentRegistry.agents[name] = AgentRegistryEntry(class_name, module_name)

deep_quoridor/src/agents/nn/pyramid_1024_dropout.py

@@ -62,6 +62,11 @@ def observation_to_tensor(self, observation):
         for key, value in observation.items():
             if isinstance(value, np.ndarray):
                 flat_obs.extend(value.flatten())
+            elif key == "player_turn":
+                if value == "player_0":
+                    value = 0
+                else:
+                    value = 1
             else:
                 flat_obs.append(value)
         return torch.FloatTensor(flat_obs).to(my_device())

deep_quoridor/src/agents/sb3_ppo.py

@@ -292,7 +292,3 @@ def forward(self, obs: dict) -> torch.Tensor:
         for v in obs.values():
             thobs = torch.cat((thobs, torch.tensor(v).flatten(start_dim=1)), dim=1)
         return thobs
-
-
-# Register the agent with the registry
-AgentRegistry.register(Agent._friendly_name(SB3PPOAgent.__name__), SB3PPOAgent)