Add baseline advisor for lr model

requirements-lock.txt

@@ -1,5 +1,5 @@
 astroid==3.3.5
-baseline_models @ git+https://git@github.com/ALLAN-DIP/baseline-models.git@ba5700a587ffab63e35a44d9ed96919a4ffee0cf
+baseline_models @ git+https://git@github.com/ALLAN-DIP/baseline-models.git@367fc4bb269e37ddcd60978e3a11cf180106679c
 bcrypt==4.2.0
 certifi==2024.7.4
 cfgv==3.4.0

requirements.txt

@@ -1,4 +1,4 @@
-baseline-models @ git+https://git@github.com/ALLAN-DIP/baseline-models.git@ba5700a587ffab63e35a44d9ed96919a4ffee0cf
+baseline-models @ git+https://git@github.com/ALLAN-DIP/baseline-models.git@367fc4bb269e37ddcd60978e3a11cf180106679c
 daidepp @ git+https://git@github.com/SHADE-AI/daidepp.git@7f6dcca60f2e6dd89c37862663be7b098cc03794
 # Pinned `diplomacy` to `main`
 diplomacy @ git+https://git@github.com/ALLAN-DIP/diplomacy.git@f4ec4517dc5ea3c3ed787d3ad483ce01f4f0e0ea

src/chiron_utils/bots/__init__.py

@@ -7,6 +7,10 @@
     KnnAdvisor as KnnAdvisor,
     KnnPlayer as KnnPlayer,
 )
+from chiron_utils.bots.lr_bot import (
+    LrAdvisor as LrAdvisor,
+    LrPlayer as LrPlayer,
+)
 from chiron_utils.bots.random_proposer_bot import (
     RandomProposerAdvisor as RandomProposerAdvisor,
     RandomProposerPlayer as RandomProposerPlayer,
@@ -15,6 +19,8 @@
 BOTS: List[Type[BaselineBot]] = [
     KnnPlayer,
     KnnAdvisor,
+    LrPlayer,
+    LrAdvisor,
     RandomProposerAdvisor,
     RandomProposerPlayer,
 ]

src/chiron_utils/bots/knn_bot.py

@@ -7,8 +7,7 @@
 import pickle
 from typing import Any, Dict, List, Sequence
 
-from baseline_models.model_code.evaluation import infer
-from baseline_models.model_code.preprocess import entry_to_vectors
+from baseline_models.model_code.predict import predict
 from diplomacy.utils import strings as diplomacy_strings
 from diplomacy.utils.constants import SuggestionType
 
@@ -19,22 +18,20 @@
 
 DEFAULT_COMM_STAGE_LENGTH = 300  # 5 minutes in seconds
 COMM_STAGE_LENGTH = int(os.environ.get("COMM_STAGE_LENGTH", DEFAULT_COMM_STAGE_LENGTH))
-MODEL_PATH = Path() / "baseline_knn_model.pkl"
-
-POWER_TO_INDEX = {
-    "AUSTRIA": 0,
-    "ENGLAND": 1,
-    "FRANCE": 2,
-    "GERMANY": 3,
-    "ITALY": 4,
-    "RUSSIA": 5,
-    "TURKEY": 6,
-}
+MODEL_PATH = Path() / "knn_models"
 
 
 @dataclass
 class KnnBot(BaselineBot, ABC):
-    """Currently a dictionary mapping phase type to a model
+    """Baseline knn model
+
+    MODEL_PATH should point to folder containing model .pkl files
+
+    Each model corresponds to a (unit, location, phase) combination
+
+    Unit types are 'A', 'F'
+
+    Location types include all possible locations on the board, such as 'LON', 'STP_SC', 'BRE', etc
 
     Phase types are 'SM', 'FM', 'WA, 'SR', 'FR', 'CD'
     """
@@ -46,37 +43,21 @@ def __post_init__(self) -> None:
             self.models: Dict[str, Any] = pickle.load(model_file)
 
     def get_orders(self) -> List[str]:
-        name = self.game.phase
-        units = self.game.map.units
-        centers = self.game.map.centers
-        homes = self.game.map.homes
-
-        influences = {}
-        for power, power_class in self.game.powers.items():
-            influences[power] = power_class.influence
-
-        vector, _, season = entry_to_vectors(
-            None,
-            include_orders=False,
-            name_data=name,
-            units_data=units,
-            centers_data=centers,
-            homes_data=homes,
-            influences_data=influences,
-        )
-        orders = infer(self.models[season], vector.reshape(1, -1))
+        orders = list()
+        state = self.game.get_state()
+        orders = predict(MODEL_PATH, state, self.power_name)
+
         logger.info("Orders to suggest: %s", orders)
 
         return orders
 
     async def gen_orders(self) -> List[str]:
         orders = self.get_orders()
-        power_orders = orders[POWER_TO_INDEX[self.power_name]]
         if self.bot_type == BotType.ADVISOR:
-            await self.suggest_orders(power_orders)
+            await self.suggest_orders(orders)
         elif self.bot_type == BotType.PLAYER:
             await self.send_orders(orders, wait=True)
-        return power_orders
+        return orders
 
     async def do_messaging_round(self, orders: Sequence[str]) -> List[str]:
         return list(orders)

src/chiron_utils/bots/lr_bot.py

@@ -0,0 +1,71 @@
+"""Abstract base classes for bots."""
+
+from abc import ABC
+from dataclasses import dataclass
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Sequence
+
+from baseline_models.model_code.predict import predict
+from diplomacy.utils import strings as diplomacy_strings
+from diplomacy.utils.constants import SuggestionType
+
+from chiron_utils.bots.baseline_bot import BaselineBot, BotType
+from chiron_utils.utils import return_logger
+
+logger = return_logger(__name__)
+
+DEFAULT_COMM_STAGE_LENGTH = 300  # 5 minutes in seconds
+COMM_STAGE_LENGTH = int(os.environ.get("COMM_STAGE_LENGTH", DEFAULT_COMM_STAGE_LENGTH))
+MODEL_PATH = Path() / "lr_models"
+
+
+@dataclass
+class LrBot(BaselineBot, ABC):
+    """Baseline lr model
+
+    MODEL_PATH should point to folder containing model .pkl files
+
+    Each model corresponds to a (unit, location, phase) combination
+
+    Unit types are 'A', 'F'
+
+    Location types include all possible locations on the board, such as 'LON', 'STP_SC', 'BRE', etc
+
+    Phase types are 'SM', 'FM', 'WA, 'SR', 'FR', 'CD'
+    """
+
+    player_type = diplomacy_strings.NO_PRESS_BOT
+
+    def get_orders(self) -> List[str]:
+        orders = list()
+        state = self.game.get_state()
+        orders = predict(MODEL_PATH, state, self.power_name)
+
+        logger.info("Orders to suggest: %s", orders)
+
+        return orders
+
+    async def gen_orders(self) -> List[str]:
+        orders = self.get_orders()
+        if self.bot_type == BotType.ADVISOR:
+            await self.suggest_orders(orders)
+        elif self.bot_type == BotType.PLAYER:
+            await self.send_orders(orders, wait=True)
+        return orders
+
+    async def do_messaging_round(self, orders: Sequence[str]) -> List[str]:
+        return list(orders)
+
+
+class LrAdvisor(LrBot):
+    """Advisor form of `LrBot`."""
+
+    bot_type = BotType.ADVISOR
+    suggestion_type = SuggestionType.MOVE
+
+
+class LrPlayer(LrBot):
+    """Player form of `LrBot`."""
+
+    bot_type = BotType.PLAYER