diff --git a/deep_quoridor/src/agents/__init__.py b/deep_quoridor/src/agents/__init__.py index 2ccccc1e..121cb4d3 100644 --- a/deep_quoridor/src/agents/__init__.py +++ b/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") diff --git a/deep_quoridor/src/agents/core/agent.py b/deep_quoridor/src/agents/core/agent.py index dab095c2..7daaa6b4 100644 --- a/deep_quoridor/src/agents/core/agent.py +++ b/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) diff --git a/deep_quoridor/src/agents/nn/pyramid_1024_dropout.py b/deep_quoridor/src/agents/nn/pyramid_1024_dropout.py index c76439ee..f957d7f8 100644 --- a/deep_quoridor/src/agents/nn/pyramid_1024_dropout.py +++ b/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()) diff --git a/deep_quoridor/src/agents/sb3_ppo.py b/deep_quoridor/src/agents/sb3_ppo.py index 3f01feb7..9e726f70 100644 --- a/deep_quoridor/src/agents/sb3_ppo.py +++ b/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)