Merge pull request #150 from AresSC2/feat/spawn-controller-archons

feat: spawn controller archons

docs/tutorials/build_runner.md

@@ -84,6 +84,11 @@ race. For each build name, ensure that there is a corresponding build name under
 be careful adding build steps that are impossible to commence. Such as adding a barracks before a 
 supply depot or a gateway before pylon.
 
+Further to this, be careful adding steps where units require morphing from other units. If you
+want banelings then it's up to the author to ensure zerglings are around. If you require an Archon
+then two templar should exist to ensure the morph is successful. The build runner will not fill
+in the gaps for you.
+
 
 ### Valid build order options
 Each item in the build order should contain a string, with the first word being the command. 

docs/tutorials/managing_production.md

@@ -30,6 +30,8 @@ def viking_tank(self) -> dict:
 ```
 Things to note:
 
+ - `SpawnController` will not produce build structures for you. Please be especially careful
+of this fact when morphing units. For example Archons, Banelings, Ravagers, Lurkers or Brood Lords.
  - The `proportion` values should add up to 1.0 (0.69 + 0.13 + 0.16 + 0.02 = 1.0)
  - Each declared unit should be given a priority, where 0 is the highest. This allows resources
 to be saved for important units.

src/ares/behaviors/macro/spawn_controller.py

@@ -12,7 +12,9 @@
 
 if TYPE_CHECKING:
     from ares import AresBot
+
 from ares.behaviors.macro import MacroBehavior
+from ares.consts import UnitRole
 from ares.managers.manager_mediator import ManagerMediator
 
 
@@ -124,9 +126,15 @@ def execute(self, ai: "AresBot", config: dict, mediator: ManagerMediator) -> boo
 
             tech_ready_for.append(unit_type_id)
 
+            trained_from: set[UnitID]
+            if unit_type_id == UnitID.ARCHON:
+                trained_from = {UnitID.HIGHTEMPLAR, UnitID.DARKTEMPLAR}
+            else:
+                trained_from = UNIT_TRAINED_FROM[unit_type_id]
+
             # get all idle build structures/units we can create this unit from
             build_structures: list[Unit] = ai.get_build_structures(
-                UNIT_TRAINED_FROM[unit_type_id],
+                trained_from,
                 unit_type_id,
                 self.__build_dict,
                 self.ignored_build_from_tags,
@@ -135,6 +143,11 @@ def execute(self, ai: "AresBot", config: dict, mediator: ManagerMediator) -> boo
             if len(build_structures) == 0:
                 continue
 
+            # archon is a special case that can't be handled generically
+            if unit_type_id == UnitID.ARCHON:
+                self._handle_archon_morph(ai, build_structures, mediator)
+                continue
+
             # prioritize spawning close to spawn target
             if self.spawn_target:
                 build_structures = cy_sorted_by_distance_to(
@@ -143,7 +156,7 @@ def execute(self, ai: "AresBot", config: dict, mediator: ManagerMediator) -> boo
 
             # can't afford unit?
             # then we might want to break out loop till we can afford
-            if not ai.can_afford(unit_type_id):
+            if not self._can_afford(ai, unit_type_id):
                 if (
                     self.freeflow_mode
                     or num_total_units < self.ignore_proportions_below_unit_count
@@ -209,25 +222,7 @@ def execute(self, ai: "AresBot", config: dict, mediator: ManagerMediator) -> boo
                 proportion_sum, 1.0
             ), f"The army comp proportions should equal 1.0, got {proportion_sum}"
 
-        did_action: bool = False
-        for unit, value in self.__build_dict.items():
-            did_action = True
-            mediator.clear_role(tag=unit.tag)
-            if value == UnitID.BANELING:
-                unit(AbilityId.MORPHTOBANELING_BANELING)
-            elif value == UnitID.RAVAGER:
-                unit(AbilityId.MORPHTORAVAGER_RAVAGER)
-            # prod building is warp gate, but we really
-            # want to spawn from psionic field
-            elif unit.type_id == UnitID.WARPGATE:
-                mediator.request_warp_in(
-                    build_from=unit, unit_type=value, target=self.spawn_target
-                )
-            else:
-                unit.train(value)
-                ai.num_larva_left -= 1
-
-        return did_action
+        return self._morph_units(ai, mediator)
 
     def _add_to_build_dict(
         self,
@@ -299,3 +294,46 @@ def _calculate_build_amount(
             maximum,
         )
         return amount, supply_cost, cost
+
+    @staticmethod
+    def _can_afford(ai: "AresBot", unit_type_id: UnitID) -> bool:
+        if unit_type_id == UnitID.ARCHON:
+            return True
+        return ai.can_afford(unit_type_id)
+
+    @staticmethod
+    def _handle_archon_morph(
+        ai: "AresBot", build_structures: list[Unit], mediator: ManagerMediator
+    ) -> None:
+        unit_role_dict: dict[UnitRole, set] = mediator.get_unit_role_dict
+        build_structures = [
+            b
+            for b in build_structures
+            if b.tag not in unit_role_dict[UnitRole.MORPHING] and b.is_ready
+        ]
+        if len(build_structures) < 2:
+            return
+
+        templar: list[Unit] = build_structures[:2]
+        ai.request_archon_morph(templar)
+
+    def _morph_units(self, ai: "AresBot", mediator: ManagerMediator) -> bool:
+        did_action: bool = False
+        for unit, value in self.__build_dict.items():
+            did_action = True
+            mediator.clear_role(tag=unit.tag)
+            if value == UnitID.BANELING:
+                unit(AbilityId.MORPHTOBANELING_BANELING)
+            elif value == UnitID.RAVAGER:
+                unit(AbilityId.MORPHTORAVAGER_RAVAGER)
+            # prod building is warp gate, but we really
+            # want to spawn from psionic field
+            elif unit.type_id == UnitID.WARPGATE:
+                mediator.request_warp_in(
+                    build_from=unit, unit_type=value, target=self.spawn_target
+                )
+            else:
+                unit.train(value)
+                ai.num_larva_left -= 1
+
+        return did_action

src/ares/build_runner/build_order_parser.py

@@ -179,20 +179,24 @@ def _generate_unit_build_step(self, unit_id: UnitID) -> Callable:
         BuildOrderStep :
             A new build step to put in a build order.
         """
+        trained_from: set[UnitID]
+        if unit_id == UnitID.ARCHON:
+            trained_from = {UnitID.DARKTEMPLAR, UnitID.HIGHTEMPLAR}
+        else:
+            trained_from = UNIT_TRAINED_FROM[unit_id]
+
+        check_supply_cost: bool = unit_id not in {UnitID.ARCHON, UnitID.BANELING}
         return lambda: BuildOrderStep(
             command=unit_id,
-            start_condition=lambda: self.ai.can_afford(unit_id)
-            and self.ai.tech_ready_for_unit(unit_id)
-            and len(
-                self.ai.get_build_structures(
-                    UNIT_TRAINED_FROM[unit_id],
-                    unit_id,
-                )
+            start_condition=lambda: (
+                self.ai.can_afford(unit_id, check_supply_cost=check_supply_cost)
+                or unit_id == UnitID.ARCHON
             )
-            > 0,
+            and self.ai.tech_ready_for_unit(unit_id)
+            and len(self.ai.get_build_structures(trained_from, unit_id)) > 0,
             # if start condition is True a train order will be issued
             # therefore it will automatically complete the step
-            end_condition=lambda: True,
+            end_condition=lambda: unit_id != UnitID.ARCHON,
         )
 
     def _generate_upgrade_build_step(self, upgrade_id: UpgradeId) -> Callable:

src/ares/build_runner/build_order_runner.py

@@ -357,6 +357,11 @@ async def do_step(self, step: BuildOrderStep) -> None:
                 # backup here just in case
                 elif isinstance(command, UpgradeId):
                     self.ai.research(command)
+                elif command == UnitID.ARCHON:
+                    army_comp: dict = {command: {"proportion": 1.0, "priority": 0}}
+                    SpawnController(army_comp, freeflow_mode=True, maximum=1).execute(
+                        self.ai, self.config, self.mediator
+                    )
 
             # end condition active, complete step
             else:

src/ares/consts.py

@@ -389,6 +389,7 @@ class UnitRole(str, Enum):
     HARASSING = "HARASSING"  # units that are harassing
     IDLE = "IDLE"  # not doing anything
     MAP_CONTROL = "MAP_CONTROL"  # units controlling the map (lings/hellions?)
+    MORPHING = "MORPHING"  # units currently morphing
     OFFENSIVE_REPAIR = "OFFENSIVE_REPAIR"  # with the main force
     OVERLORD_HUNTER = "OVERLORD_HUNTER"  # units looking for overlords
     PERSISTENT_BUILDER = "PERSISTENT_BUILDER"  # does not get reassigned automatically

src/ares/custom_bot_ai.py

@@ -251,6 +251,17 @@ async def _give_units_same_order(
                 )
             )
 
+    async def _do_archon_morph(self, templar: list[Unit]) -> None:
+        command = raw_pb.ActionRawUnitCommand(
+            ability_id=AbilityId.MORPH_ARCHON.value,
+            unit_tags=[templar[0].tag, templar[1].tag],
+            queue_command=False,
+        )
+        action = raw_pb.ActionRaw(unit_command=command)
+        await self.client._execute(
+            action=sc_pb.RequestAction(actions=[sc_pb.Action(action_raw=action)])
+        )
+
     async def unload_by_tag(self, container: Unit, unit_tag: int) -> None:
         """Unload a unit from a container based on its tag. Thanks, Sasha!"""
         index: int = 0

src/ares/dicts/unit_tech_requirement.py

@@ -3,6 +3,7 @@
 UNIT_TECH_REQUIREMENT: dict[UnitID, list[UnitID]] = dict(
     {
         UnitID.ADEPT: [UnitID.GATEWAY, UnitID.CYBERNETICSCORE],
+        UnitID.ARCHON: [],
         UnitID.BANELING: [UnitID.BANELINGNEST],
         UnitID.BANSHEE: [UnitID.STARPORT, UnitID.STARPORTTECHLAB],
         UnitID.BARRACKS: [UnitID.SUPPLYDEPOT],

src/ares/main.py

@@ -110,6 +110,7 @@ def __init__(self, game_step_override: Optional[int] = None):
             tuple[AbilityId, set[int], Optional[Union[Unit, Point2]]]
         ] = []
         self._drop_unload_actions: list[tuple[int, int]] = []
+        self._archon_morph_actions: list[list] = []
 
         self.arcade_mode: bool = False
 
@@ -124,6 +125,9 @@ def give_same_action(
     def do_unload_container(self, container_tag: int, index: int = 0) -> None:
         self._drop_unload_actions.append((container_tag, index))
 
+    def request_archon_morph(self, templar: list[Unit]) -> None:
+        self._archon_morph_actions.append(templar)
+
     # noinspection PyFinal
     def _prepare_units(self):
         """Tweak of _prepare_units to include memory units in cached distances and some
@@ -378,6 +382,8 @@ async def _after_step(self) -> int:
             await self._give_units_same_order(
                 same_order[0], same_order[1], same_order[2]
             )
+        for archon_morph_action in self._archon_morph_actions:
+            await self._do_archon_morph(archon_morph_action)
         self.manager_hub.path_manager.reset_grids(self.actual_iteration)
         await self.manager_hub.placement_manager.do_warp_ins()
         return await super(AresBot, self)._after_step()
@@ -471,6 +477,11 @@ async def on_unit_created(self, unit: Unit) -> None:
             and unit.type_id in GATEWAY_UNITS
         ):
             self.build_order_runner.set_step_started(True)
+        if (
+            not self.build_order_runner.build_completed
+            and unit.type_id == UnitID.ARCHON
+        ):
+            self.build_order_runner.set_step_complete(UnitID.ARCHON)
         await self.manager_hub.on_unit_created(unit)
 
     async def on_unit_destroyed(self, unit_tag: int) -> None:
@@ -795,6 +806,7 @@ def _reset_variables(self) -> None:
 
         self._drop_unload_actions = []
         self._same_order_actions = []
+        self._archon_morph_actions = []
 
     def _should_add_unit(self, unit: RawUnit) -> bool:
         """Whether the given unit should be tracked.
@@ -882,7 +894,7 @@ def get_build_structures(
         structures_dict: dict[UnitID:Units] = self.mediator.get_own_structures_dict
         own_army_dict: dict[UnitID:Units] = self.mediator.get_own_army_dict
         build_from_dict: dict[UnitID:Units] = structures_dict
-        if self.race == Race.Zerg:
+        if self.race != Race.Terran:
             build_from_dict: dict[UnitID:Units] = {
                 **structures_dict,
                 **own_army_dict,
@@ -905,9 +917,9 @@ def get_build_structures(
                     if AbilityId.WARPGATETRAIN_ZEALOT in b.abilities
                 ]
 
-            requires_techlab: bool = TRAIN_INFO[structure_type][unit_type].get(
-                "requires_techlab", False
-            )
+            requires_techlab: bool = self.race == Race.Terran and TRAIN_INFO[
+                structure_type
+            ][unit_type].get("requires_techlab", False)
             if not requires_techlab:
                 build_from_tags.extend(
                     [
@@ -953,7 +965,10 @@ def get_build_structures(
         # limit to powered structures
         if self.race == Race.Protoss:
             build_structures = [
-                s for s in build_structures if s.is_powered or s.type_id == UnitID.NEXUS
+                s
+                for s in build_structures
+                if s.is_powered
+                or s.type_id in {UnitID.NEXUS, UnitID.DARKTEMPLAR, UnitID.HIGHTEMPLAR}
             ]
 
         return build_structures