[WIP] Build order support

.gitignore

@@ -13,3 +13,6 @@ maps/
 mini_games/
 
 *.SC2Replay
+
+.idea
+*.iml

examples/terran_marine_rush_build_order.py

@@ -0,0 +1,64 @@
+import sc2
+from sc2 import run_game, maps, Race, Difficulty
+from sc2.build_orders.build_order import BuildOrder, train_unit
+from sc2.build_orders.commands import construct, expand, add_supply
+from sc2.constants import *
+from sc2.player import Bot, Computer
+from sc2.state_conditions.conditions import all_of, supply_at_least, minerals_at_least, unit_count
+
+
+def first_barracks(bot):
+    return bot.units(UnitTypeId.BARRACKS).first
+
+
+class TerranBuildOrderBot(sc2.BotAI):
+    def __init__(self):
+        build_order = [
+            (supply_at_least(13), add_supply(prioritize=True)),
+            (supply_at_least(15), construct(UnitTypeId.BARRACKS, prioritize=True)),
+            (supply_at_least(16), construct(UnitTypeId.BARRACKS, prioritize=True)),
+            (supply_at_least(16), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(17), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(18), construct(UnitTypeId.BARRACKS, prioritize=True)),
+            (supply_at_least(18), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(19), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(20), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(21), add_supply(prioritize=True)),
+            (supply_at_least(21), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(22), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(23), construct(UnitTypeId.BARRACKS, prioritize=True)),
+            (supply_at_least(23), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(24), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(25), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(26), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS)),
+            (supply_at_least(27), construct(UnitTypeId.BARRACKS, prioritize=True)),
+            (all_of(supply_at_least(27), unit_count(UnitTypeId.BARRACKS, 5, include_pending=True)), add_supply(prioritize=True)),
+            (unit_count(UnitTypeId.BARRACKS, 5), train_unit(UnitTypeId.MARINE, on_building=UnitTypeId.BARRACKS, repeatable=True))
+        ]
+        self.attack = False
+        self.build_order = BuildOrder(self, build_order, worker_count=16)
+
+    async def on_step(self, iteration):
+        await self.distribute_workers()
+        await self.build_order.execute_build()
+
+        if self.units(UnitTypeId.MARINE).amount >= 15 or self.attack:
+            self.attack = True
+            for unit in self.units(UnitTypeId.MARINE).idle:
+                await self.do(unit.attack(self.enemy_start_locations[0]))
+                if self.known_enemy_structures.exists:
+                    enemy = self.known_enemy_structures.first
+                    await self.do(unit.attack(enemy.position.to2, queue=True))
+            return
+
+
+def main():
+    run_game(maps.get("Abyssal Reef LE"), [
+        Bot(Race.Terran, TerranBuildOrderBot()),
+        # Bot(Race.Terran, TerranBuildOrderBot()),
+        Computer(Race.Random, Difficulty.Harder)
+    ], realtime=False)
+
+if __name__ == '__main__':
+    main()
+

examples/threebase_voidray.py

@@ -7,7 +7,7 @@
 from sc2.player import Bot, Computer
 
 class ThreebaseVoidrayBot(sc2.BotAI):
-    def select_target(self, state):
+    def select_target(self):
         if self.known_enemy_structures.exists:
             return random.choice(self.known_enemy_structures)
 

examples/zerg_rush_build_order.py

@@ -0,0 +1,72 @@
+import sc2
+from sc2 import Race, Difficulty
+from sc2.build_orders.build_order import train_unit, BuildOrder, morph
+from sc2.build_orders.commands import add_gas, construct, expand, add_supply
+from sc2.constants import *
+from sc2.player import Bot, Computer
+from sc2.state_conditions.conditions import supply_at_least, all_of, unit_count, gas_less_than, unit_count_less_than, \
+    unit_count_at_least
+
+
+def research(building, upgrade):
+    async def research_spec(bot):
+        sp = bot.units(building).ready
+        if sp.exists and bot.can_afford(upgrade) and not bot.already_pending(upgrade):
+            await bot.do(sp.first(upgrade))
+
+    return research_spec
+
+class ZergRushBot(sc2.BotAI):
+
+    def __init__(self):
+        self.attack = False
+        self.mboost_started = False
+        build_order = [
+            (all_of(supply_at_least(13), unit_count(UnitTypeId.OVERLORD, 1, include_pending=True)), add_supply(prioritize=True)),
+            (all_of(supply_at_least(17), unit_count(UnitTypeId.EXTRACTOR, 0, include_pending=True)), add_gas()),
+            (all_of(supply_at_least(17), unit_count(UnitTypeId.SPAWNINGPOOL, 0, include_pending=True)), construct(UnitTypeId.SPAWNINGPOOL)),
+            (all_of(supply_at_least(17), unit_count(UnitTypeId.HATCHERY, 1, include_pending=True)), expand()),
+            (supply_at_least(18), morph(UnitTypeId.ZERGLING)),
+            (supply_at_least(19), train_unit(UnitTypeId.QUEEN, on_building=UnitTypeId.HATCHERY, prioritize=True)),
+            (all_of(supply_at_least(21), unit_count(UnitTypeId.OVERLORD, 2, include_pending=True)), add_supply(prioritize=True)),
+            (all_of(supply_at_least(21), unit_count(UnitTypeId.ROACHWARREN, 0, include_pending=True)), construct(UnitTypeId.ROACHWARREN)),
+            (all_of(supply_at_least(20), unit_count(UnitTypeId.OVERLORD, 3, include_pending=True)), add_supply(prioritize=True)),
+            (unit_count_at_least(UnitTypeId.ROACH, 7), morph(UnitTypeId.ZERGLING, repeatable=True)),
+            (unit_count(UnitTypeId.ROACHWARREN, 1), morph(UnitTypeId.ROACH, repeatable=True))
+        ]
+
+        self.build_order = BuildOrder(self, build_order, worker_count=35)
+
+    async def on_step(self, iteration):
+        await self.distribute_workers()
+
+        if self.vespene >= 100:
+            sp = self.units(UnitTypeId.SPAWNINGPOOL).ready
+            if sp.exists and self.minerals >= 100 and not self.mboost_started:
+                await self.do(sp.first(AbilityId.RESEARCH_ZERGLINGMETABOLICBOOST))
+                self.mboost_started = True
+
+        await self.build_order.execute_build()
+
+        for queen in self.units(UnitTypeId.QUEEN).idle:
+            if queen.energy >= 25:  # Hard coded, since this is not (yet) available
+                hatchery = self.townhalls.closest_to(queen.position.to2)
+                await self.do(queen(AbilityId.EFFECT_INJECTLARVA, hatchery))
+
+        if (self.units(UnitTypeId.ROACH).amount >= 7 and self.units(UnitTypeId.ZERGLING).amount >= 10) or self.attack:
+            self.attack = True
+            for unit in self.units(UnitTypeId.ZERGLING).ready.idle | self.units(UnitTypeId.ROACH).ready.idle:
+                await self.do(unit.attack(self.enemy_start_locations[0]))
+                if self.known_enemy_structures.exists:
+                    enemy = self.known_enemy_structures.first
+                    await self.do(unit.attack(enemy.position.to2, queue=True))
+            return
+
+def main():
+    sc2.run_game(sc2.maps.get("Abyssal Reef LE"), [
+        Bot(Race.Zerg, ZergRushBot()),
+        Computer(Race.Random, Difficulty.Hard)
+    ], realtime=False)
+
+if __name__ == '__main__':
+    main()

sc2/bot_ai.py

@@ -7,7 +7,8 @@
 from .constants import EGG
 
 from .position import Point2, Point3
-from .data import Race, ActionResult, Attribute, race_worker, race_townhalls, race_gas
+from .data import Race, ActionResult, Attribute, race_worker, race_townhalls, race_gas, race_supply, \
+    race_basic_townhalls
 from .unit import Unit
 from .cache import property_cache_forever
 from .game_data import AbilityData, Cost
@@ -28,6 +29,11 @@ def _prepare_start(self, client, player_id, game_info, game_data):
         self.player_id = player_id
         self.race = Race(self._game_info.player_races[self.player_id])
 
+        self.worker_type = race_worker[self.race]
+        self.basic_townhall_type = race_basic_townhalls[self.race]
+        self.geyser_type = race_gas[self.race]
+        self.supply_type = race_supply[self.race]
+
     @property
     def game_info(self):
         return self._game_info
@@ -79,7 +85,7 @@ async def get_available_abilities(self, unit):
 
     async def expand_now(self, building=None, max_distance=10, location=None):
         if not building:
-            building = self.townhalls.first.type_id
+            building = self.basic_townhall_type
 
         assert isinstance(building, UnitTypeId)
 
@@ -282,7 +288,7 @@ async def chat_send(self, message):
     def _prepare_step(self, state):
         self.state = state
         self.units = state.units.owned
-        self.workers = self.units(race_worker[self.race])
+        self.workers = self.units(self.worker_type)
         self.townhalls = self.units(race_townhalls[self.race])
         self.geysers = self.units(race_gas[self.race])
 

sc2/build_orders/build_order.py

@@ -0,0 +1,41 @@
+from sc2 import Race, race_worker, ActionResult, race_townhalls
+from sc2.build_orders.commands import add_supply, morph, train_unit
+from sc2.state_conditions.conditions import always_true
+
+
+class BuildOrder(object):
+    def __init__(self, bot, build, worker_count=0, auto_add_supply=True):
+        self.build = build
+        self.bot = bot
+        self.worker_count = worker_count
+        self.auto_add_supply = auto_add_supply
+
+    async def execute_build(self):
+        bot = self.bot
+        if bot.supply_left <= ((bot.supply_cap+50) / 50) and not bot.already_pending(bot.supply_type) \
+                and self.auto_add_supply:
+            return await add_supply().execute(bot)
+
+        for index, item in enumerate(self.build):
+            condition, command = item
+            condition = item[0] if item[0] else always_true
+            if condition(bot) and not command.is_done:
+                e = await command.execute(bot)
+                if command.is_done:
+                    return e
+                else:
+                    # Save up to be able to do this command and hold worker creation.
+                    if command.is_priority and e == ActionResult.NotEnoughMinerals:
+                        return e
+
+                    if e == ActionResult.NotEnoughFood and self.auto_add_supply \
+                            and not bot.already_pending(bot.supply_type):
+                        return await add_supply().execute(bot)
+                    continue
+
+        if bot.workers.amount < self.worker_count:
+            if bot.race == Race.Zerg:
+                return await morph(race_worker[Race.Zerg]).execute(bot)
+            else:
+                return await train_unit(race_worker[bot.race], race_townhalls[self.bot.race]).execute(bot)
+        return None

sc2/build_orders/commands.py

@@ -0,0 +1,122 @@
+from sc2 import ActionResult, Race
+from sc2.constants import *
+
+
+class Command(object):
+    def __init__(self, action, repeatable=False, priority=False):
+        self.action = action
+        self.is_done = False
+        self.is_repeatable = repeatable
+        self.is_priority = priority
+
+    async def execute(self, bot):
+        e = await self.action(bot)
+        if not e and not self.is_repeatable:
+            self.is_done = True
+
+        return e
+
+    def allow_repeat(self):
+        self.is_repeatable = True
+        self.is_done = False
+        return self
+
+
+def expand(prioritize=False, repeatable=True):
+    async def do_expand(bot):
+        building = bot.basic_townhall_type
+        can_afford = bot.can_afford(building)
+        if can_afford:
+            return await bot.expand_now(building=building)
+        else:
+            return can_afford.action_result
+
+    return Command(do_expand, priority=prioritize, repeatable=repeatable)
+
+
+def train_unit(unit, on_building, prioritize=False, repeatable=False):
+    async def do_train(bot):
+        buildings = bot.units(on_building).ready.noqueue
+        if buildings.exists:
+            selected = buildings.first
+            can_afford = bot.can_afford(unit)
+            if can_afford:
+                print("Training {}".format(unit))
+                return await bot.do(selected.train(unit))
+            else:
+                return can_afford.action_result
+        else:
+            return ActionResult.Error
+
+    return Command(do_train, priority=prioritize, repeatable=repeatable)
+
+
+def morph(unit, prioritize=False, repeatable=False):
+    async def do_morph(bot):
+        larvae = bot.units(UnitTypeId.LARVA)
+        if larvae.exists:
+            selected = larvae.first
+            can_afford = bot.can_afford(unit)
+            if can_afford:
+                print("Morph {}".format(unit))
+                return await bot.do(selected.train(unit))
+            else:
+                return can_afford.action_result
+        else:
+            return ActionResult.Error
+
+    return Command(do_morph, priority=prioritize, repeatable=repeatable)
+
+
+def construct(building, placement=None, prioritize=True, repeatable=False):
+    async def do_build(bot):
+
+        if not placement:
+            location = bot.townhalls.first.position.towards(bot.game_info.map_center, 5)
+        else:
+            location = placement
+
+        can_afford = bot.can_afford(building)
+        if can_afford:
+            print("Building {}".format(building))
+            return await bot.build(building, near=location)
+        else:
+            return can_afford.action_result
+
+    return Command(do_build, priority=prioritize, repeatable=repeatable)
+
+
+def add_supply(prioritize=True, repeatable=False):
+    async def supply_spec(bot):
+        can_afford = bot.can_afford(bot.supply_type)
+        if can_afford:
+            if bot.race == Race.Zerg:
+                return await morph(bot.supply_type).execute(bot)
+            else:
+                return await construct(bot.supply_type).execute(bot)
+        else:
+            return can_afford.action_result
+
+    return Command(supply_spec, priority=prioritize, repeatable=repeatable)
+
+
+def add_gas(prioritize=True, repeatable=False):
+    async def do_add_gas(bot):
+        can_afford = bot.can_afford(bot.geyser_type)
+        if not can_afford:
+            return can_afford.action_result
+
+        owned_expansions = bot.owned_expansions
+        for location, th in owned_expansions.items():
+            vgs = bot.state.vespene_geyser.closer_than(20.0, th)
+            for vg in vgs:
+                worker = bot.select_build_worker(vg.position)
+                if worker is None:
+                    break
+
+                if not bot.units(bot.geyser_type).closer_than(1.0, vg).exists:
+                    return await bot.do(worker.build(bot.geyser_type, vg))
+
+        return ActionResult.Error
+
+    return Command(do_add_gas, priority=prioritize, repeatable=repeatable)

sc2/data.py

@@ -11,7 +11,7 @@
 from .ids.unit_typeid import COMMANDCENTER, ORBITALCOMMAND, PLANETARYFORTRESS
 from .ids.unit_typeid import HATCHERY, LAIR, HIVE
 from .ids.unit_typeid import ASSIMILATOR, REFINERY, EXTRACTOR
-
+from .ids.unit_typeid import PYLON, OVERLORD, SUPPLYDEPOT
 from .ids.ability_id import (
     GATEWAYTRAIN_ZEALOT,
     GATEWAYTRAIN_STALKER,
@@ -47,12 +47,24 @@
 
 ActionResult = enum.Enum("ActionResult", error_pb.ActionResult.items())
 
+race_supply = {
+    Race.Protoss: PYLON,
+    Race.Terran: SUPPLYDEPOT,
+    Race.Zerg: OVERLORD
+}
+
 race_worker = {
     Race.Protoss: PROBE,
     Race.Terran: SCV,
     Race.Zerg: DRONE
 }
 
+race_basic_townhalls = {
+    Race.Protoss: NEXUS,
+    Race.Terran: COMMANDCENTER,
+    Race.Zerg: HATCHERY
+}
+
 race_townhalls = {
     Race.Protoss: {NEXUS},
     Race.Terran: {COMMANDCENTER, ORBITALCOMMAND, PLANETARYFORTRESS},