Add zergling engagement and roach pylon attack behaviors

bot/main.py

@@ -6,10 +6,12 @@
     AMove,
     PathUnitToTarget,
     StutterUnitBack,
+    AttackTarget,
 )
 
 from sc2.ids.unit_typeid import UnitTypeId
 from sc2.units import Units
+from sc2.unit import Unit
 
 import numpy as np
 from sc2.units import Units
@@ -37,7 +39,7 @@ async def on_step(self, iteration: int):
 
         #define targets and grid
         enemy_units = self.enemy_units
-        ground_grid = np.ndarray = self.game_info.pathing_grid
+        ground_grid = self.game_info.pathing_grid
 
       #call the engagement for zergling and pylon attack for roaches
         self.do_zergling_engagement(zerglings, enemy_units, ground_grid)
@@ -46,42 +48,60 @@ async def on_step(self, iteration: int):
 
     def do_zergling_engagement(
         self,
-        zergling: Units,
+        zerglings: Units,
         enemies: Units,
+        grid: np.ndarray,
     ) -> None:
-        """Engage enemy units with zergling
+        """Engage enemy units with zerglings and use stutter step behavior.
 
         Parameters
         ----------
-        units :
-            zergling to engage with
-        enemies :
-            Enemy units to engage
-        grid :
-            Grid of the ground
+        zerglings : Units
+            Zerglings to engage with.
+        enemies : Units
+            Enemy units to engage.
+        grid : np.ndarray
+            The ground grid for pathing information.
         """
-        zergling_maneuver: CombatManeuver = CombatManeuver()
-        # engage enemy and retreat if needed
-        zergling_maneuver.add(StutterUnitBack(unit=zergling, target=enemies))
-        zergling_maneuver.add(AMove(unit=zergling, target=enemies))
-        self.register_behavior(zergling_maneuver)
+        for zergling in zerglings:
+            closest_enemy: Unit = cy_closest_to(zergling, enemies)
+            if closest_enemy:
+                zergling_maneuver = CombatManeuver()
+                zergling_maneuver.add(StutterUnitBack(zergling, closest_enemy, grid=grid))
+                zergling_maneuver.add(AMove(zergling, self.enemy_start_locations[0]))
+                self.register_behavior(zergling_maneuver)
 
     def do_roach_pylon_attack(
         self,
         roaches: Units,
         grid: np.ndarray,
     ) -> None:
-        """Attack move towards the enemy pylon and when it sees it prioritize it"""
-        roach_maneuver: CombatManeuver = CombatManeuver()
-        # move towards the enemy and once it sees the pylon prioritize it
-        roach_maneuver.add(
-            PathUnitToTarget(
-                unit=roaches,
-                grid=grid,
-                target=self.enemy_start_locations[0],
-                success_at_distance=5.0,
-            )
-        )
-        # roach_maneuver.add(AttackTarget(unit=roaches, target=self.enemy_structures.first))
+        """Direct roaches to move towards the enemy start location and attack the pylon once in range.
+
+        Parameters
+        ----------
+        roaches : Units
+            Roaches to perform the attack.
+        grid : np.ndarray
+            The ground grid for pathing information.
+        """
+        enemy_start_location = self.enemy_start_locations[0]
+        enemy_pylon = self._find_enemy_pylon()
+
+        for roach in roaches:
+            roach_maneuver = CombatManeuver()
+            # Path to the enemy start location
+            roach_maneuver.add(PathUnitToTarget(roach, grid, enemy_start_location, success_at_distance=5.0))
+
+            # If a pylon is found, attack it
+            if enemy_pylon:
+                roach_maneuver.add(AttackTarget(roach, enemy_pylon))
+
+            self.register_behavior(roach_maneuver)
 
-        self.register_behavior(roach_maneuver)
+    def _find_enemy_pylon(self) -> Unit:
+        """Find the enemy pylon to target."""
+        pylons = self.enemy_structures(UnitTypeId.PYLON)
+        if pylons:
+            return pylons.closest_to(self.enemy_start_locations[0])
+        return None