def select_idle_workers(self, game_state, number=None):
# TODO: Select according to race
worker_id = "SCV"
idle_workers = game_state.player_units.filter(name=worker_id, orders=[])
# Returning workers
harvest_return_abilities = [
ability.value
for ability in AbilityId
if ability.name.startswith("HARVEST_RETURN")
]
for ability_id in harvest_return_abilities:
idle_workers += game_state.player_units.filter(
name=worker_id, orders__ability_id=ability_id
)
# Gathering workers
harvest_return_abilities = [
ability.value
for ability in AbilityId
if ability.name.startswith("HARVEST_GATHER")
]
for ability_id in harvest_return_abilities:
idle_workers += game_state.player_units.filter(
name=worker_id, orders__ability_id=ability_id
)
# Return manager with idle workers
if number is not None:
return UnitManager(idle_workers[:number])
else:
return UnitManager(idle_workers)
def determine_action_state(self, game_state, units_queue, upgrades_queue):
available_builders = UnitManager(
get_available_building_unit(self.unit_id, game_state))
builder_orders = len(
list(
itertools.chain(
*available_builders.values('orders', flat_list=True))))
no_space_in_queue = builder_orders >= 5 * len(available_builders)
if available_builders and no_space_in_queue:
return Action.MISSING_SPACE_IN_QUEUE, []
else:
return super(Train,
self).determine_action_state(game_state, units_queue,
upgrades_queue)
def get_closing_enemies(game_state):
from game_data.units import UnitManager
town_halls = game_state.player_units.filter(unit_type__in=[
UnitTypeIds.COMMANDCENTER.value,
UnitTypeIds.ORBITALCOMMAND.value,
UnitTypeIds.PLANETARYFORTRESS.value,
])
min_distance = 30
dangerous_units = UnitManager([])
for th in town_halls:
e_units = game_state.enemy_units \
.add_calculated_values(distance_to={"unit": th}) \
.filter(last_distance_to__lte=min_distance)
dangerous_units += e_units
return dangerous_units.values('tag', flat_list=True)
def prepare_unit_composition(self, composition, distance_args, game_state):
base_manager = UnitManager([])
for unit_type, amount in composition.items():
units = game_state.player_units.filter(unit_type=unit_type)
if distance_args:
units = units.add_calculated_values(distance_to=distance_args)
base_manager += units[:amount]
return base_manager
def get_action_units(self, game_state, target):
distance_args = None
if target.get('target_unit'):
distance_args = {"unit": target.get('target_unit')}
elif target.get('target_point'):
distance_args = {"pos": target.get('target_point')}
# Get units performing the action
base_manager = UnitManager([])
composition = self.unit_group.get('composition')
query_params = self.unit_group.get('query')
if query_params:
base_manager = game_state.player_units.filter(**query_params)
if distance_args:
base_manager = base_manager.add_calculated_values(
distance_to=distance_args)
elif composition:
base_manager = self.prepare_unit_composition(
composition, distance_args, game_state)
return base_manager
async def process_order(self, ws, game_state):
if self.unit_index:
index = self.unit_index
units = UnitManager(
game_state.player_units.filter(
**self.unit_filters)[index:index + 1])
else:
units = game_state.player_units.filter(**self.unit_filters)
target_unit = None
target_point = None
if self.target:
if self.target.get("unit"):
# Select unit manager given the target's alliance
if self.target['unit'].get("alliance"):
if self.target['unit']["alliance"] == Alliance.Value(
"Enemy"):
unit_manager = game_state.enemy_units
else:
unit_manager = game_state.player_units
else:
unit_manager = game_state.player_units
unit = unit_manager.filter(
**self.target['unit']['filter_params'])[self.target['unit']
['index']]
if self.target.get("diff"):
diff = self.target["diff"]
target_point = (unit.get_attribute("pos").x + diff[0],
unit.get_attribute("pos").y + diff[1])
elif self.target['unit'].get("pos"):
target_point = (unit.get_attribute("pos").x,
unit.get_attribute("pos").y)
else:
target_unit = unit
elif self.target.get("point"):
target_point = self.target['point']
if self.ability_id in BUILD_ABILITIES:
target_point = await find_placement(ws, self.ability_id,
target_point)
result = await units.give_order(ws,
ability_id=self.ability_id,
target_point=target_point,
target_unit=target_unit)
return result
def distribute_workers_throughout_resources(
self,
workers: UnitManager,
resources_missing_workers: List[Tuple[Unit, int]],
) -> Dict[Unit, UnitManager]:
worker_distribution = {}
taken_workers = []
for resource, missing_workers in resources_missing_workers:
# Filter workers to take
workers = workers.filter(
mode=UnitManager.EXCLUDE_AND_MODE,
tag__in=taken_workers).add_calculated_values(distance_to={
'unit': resource
}).sort_by('last_distance_to')
# Take workers
workers_to_resource = workers[:missing_workers]
worker_distribution[resource] = workers_to_resource
taken_workers = workers_to_resource.values('tag', flat_list=True)
return worker_distribution
def get_action_units(self, game_state, target):
target_unit = target.get('target_unit')
missing_workers = self.ideal_workers - self.assigned_workers
# If target unit is None or no workers required, return empty list of workers
if target_unit is None or missing_workers <= 0:
return []
workers = UnitManager([])
composition = self.unit_group.get('composition')
query_params = self.unit_group.get('query')
if query_params:
type_workers = game_state.player_units.filter(
**query_params).add_calculated_values(
distance_to={
"unit": target_unit
},
unit_availability={},
).sort_by('last_unit_availability', 'last_distance_to')
assignable_workers = missing_workers
workers += type_workers[:assignable_workers]
missing_workers -= assignable_workers
elif composition:
for worker_type, worker_amount in composition.items():
type_workers = game_state.player_units.filter(
unit_type=worker_type).add_calculated_values(
distance_to={
"unit": target_unit
},
unit_availability={},
).sort_by('last_unit_availability', 'last_distance_to')
assignable_workers = min(missing_workers, worker_amount)
workers += type_workers[:assignable_workers]
missing_workers -= assignable_workers
return workers
def __init__(self, observation, game_data, game_info, actions=()):
# Set game info
self.game_info = game_info.game_info
# Set player data
self.player_info = PlayerCommon(
observation.player_common, observation.raw_data.player, game_data
)
# Set game units
self.all_units = UnitManager(
[Unit(unit, game_data) for unit in observation.raw_data.units]
)
self.player_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Self")
]
)
self.allied_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Ally")
]
)
self.enemy_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Enemy")
]
)
self.neutral_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Neutral")
]
)
self.all_units = UnitManager(
[Unit(unit, game_data) for unit in observation.raw_data.units]
)
self.enemy_currently_seeing_units = self.enemy_units.filter(display=1)
self.enemy_snapshot = self.enemy_units.filter(display=2)
# Effects and events data
self.game_event = observation.raw_data.event
self.game_effects = observation.raw_data.effects
# Observation's game loop
self.game_loop = observation.game_loop
self.game_data = game_data
self.actions = actions
self.parsed_actions = []
for action in self.actions:
ac = action.action_raw
if not ac.unit_command:
continue
else:
identifier = ac.unit_command.ability_id
types_grouped = {}
if ac.unit_command.target_unit_tag:
tag = ac.unit_command.target_unit_tag
unit = self.all_units.filter(tag=tag)
if unit:
position = unit[0].get_attribute("pos")
if unit[0].proto_unit_data:
target_unit_type = unit[0].proto_unit_data.unit_id
target_unit_name = unit[0].proto_unit_data.name
targetx = position.x
targety = position.y
else:
continue
else:
targetx, targety = (
ac.unit_command.target_world_space_pos.x,
ac.unit_command.target_world_space_pos.y,
)
target_unit_type = None
target_unit_name = None
for tag in list(ac.unit_command.unit_tags):
if self.player_units.filter(tag=tag):
unitType = self.player_units.filter(tag=tag)[0].proto_unit_data
if unitType:
order = types_grouped.get(unitType.unit_id, 0)
order += 1
types_grouped[unitType.unit_id] = order
if types_grouped.items():
for key in types_grouped.keys():
item = {
"unit_type_id": key,
"amount": types_grouped[key],
"action_id": identifier,
"x": targetx,
"y": targety,
"target_type": target_unit_type,
}
self.parsed_actions.append(item)
self.visibility_data = observation.raw_data.map_state.visibility
self.discovery_percentage = sum(
filter(lambda x: x >= 1, self.visibility_data.data)
) / len(self.visibility_data.data)
self.creep_data = observation.raw_data.map_state.creep
self.creep_percentage = sum(
filter(lambda x: x >= 1, self.creep_data.data)
) / len(self.creep_data.data)
class DecodedObservation:
def __init__(self, observation, game_data, game_info, actions=()):
# Set game info
self.game_info = game_info.game_info
# Set player data
self.player_info = PlayerCommon(
observation.player_common, observation.raw_data.player, game_data
)
# Set game units
self.all_units = UnitManager(
[Unit(unit, game_data) for unit in observation.raw_data.units]
)
self.player_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Self")
]
)
self.allied_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Ally")
]
)
self.enemy_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Enemy")
]
)
self.neutral_units = UnitManager(
[
Unit(unit, game_data)
for unit in observation.raw_data.units
if unit.alliance == api_data.Alliance.Value("Neutral")
]
)
self.all_units = UnitManager(
[Unit(unit, game_data) for unit in observation.raw_data.units]
)
self.enemy_currently_seeing_units = self.enemy_units.filter(display=1)
self.enemy_snapshot = self.enemy_units.filter(display=2)
# Effects and events data
self.game_event = observation.raw_data.event
self.game_effects = observation.raw_data.effects
# Observation's game loop
self.game_loop = observation.game_loop
self.game_data = game_data
self.actions = actions
self.parsed_actions = []
for action in self.actions:
ac = action.action_raw
if not ac.unit_command:
continue
else:
identifier = ac.unit_command.ability_id
types_grouped = {}
if ac.unit_command.target_unit_tag:
tag = ac.unit_command.target_unit_tag
unit = self.all_units.filter(tag=tag)
if unit:
position = unit[0].get_attribute("pos")
if unit[0].proto_unit_data:
target_unit_type = unit[0].proto_unit_data.unit_id
target_unit_name = unit[0].proto_unit_data.name
targetx = position.x
targety = position.y
else:
continue
else:
targetx, targety = (
ac.unit_command.target_world_space_pos.x,
ac.unit_command.target_world_space_pos.y,
)
target_unit_type = None
target_unit_name = None
for tag in list(ac.unit_command.unit_tags):
if self.player_units.filter(tag=tag):
unitType = self.player_units.filter(tag=tag)[0].proto_unit_data
if unitType:
order = types_grouped.get(unitType.unit_id, 0)
order += 1
types_grouped[unitType.unit_id] = order
if types_grouped.items():
for key in types_grouped.keys():
item = {
"unit_type_id": key,
"amount": types_grouped[key],
"action_id": identifier,
"x": targetx,
"y": targety,
"target_type": target_unit_type,
}
self.parsed_actions.append(item)
self.visibility_data = observation.raw_data.map_state.visibility
self.discovery_percentage = sum(
filter(lambda x: x >= 1, self.visibility_data.data)
) / len(self.visibility_data.data)
self.creep_data = observation.raw_data.map_state.creep
self.creep_percentage = sum(
filter(lambda x: x >= 1, self.creep_data.data)
) / len(self.creep_data.data)
def get_visibility_map(self):
image_data = map(lambda x: 100 * x, self.visibility_data.data)
image = Image.frombytes(
"L",
(self.visibility_data.size.x, self.visibility_data.size.y),
bytes(image_data),
"raw",
)
return image
def get_creep_map(self):
image_data = map(lambda x: 100 * x, self.creep_data.data)
image = Image.frombytes(
"L",
(self.creep_data.size.x, self.creep_data.size.y),
bytes(image_data),
"raw",
)
return image
def to_case(self, replay_info):
json_dict = {
"minerals": self.player_info.minerals,
"vespene": self.player_info.vespene,
"game_loop": self.game_loop,
"food_cap": self.player_info.food_cap,
"food_used": self.player_info.food_used,
"food_army": self.player_info.food_army,
"food_workers": self.player_info.food_workers,
"idle_worker_count": self.player_info.idle_worker_count,
"visibility_percentage": round(self.discovery_percentage, 2),
"creep_percentage": round(self.creep_percentage, 2),
"army_count": self.player_info.army_count,
"warp_gate_count": self.player_info.warp_gate_count,
"units": self.count_units(self.player_units),
"enemy_units": self.count_units(self.enemy_units),
"upgrades": list(map(lambda x: x.upgrade_id, self.player_info.upgrades)),
"visible_enemy_units": self.count_units(self.enemy_currently_seeing_units),
"known_invisible_enemy_units": self.count_units(self.enemy_snapshot),
"map": replay_info["map"],
"actions": self.parsed_actions,
}
return json_dict
def count_units(self, units):
type_amount = {}
for unit in units:
val = type_amount.get(unit.proto_unit.unit_type, 0)
val += 1
type_amount[unit.proto_unit.unit_type] = val
result = []
for key in type_amount.keys():
item = {"unit_type_ id": key, "amount": type_amount[key]}
result.append(item)
return result
def terrain_height(self, position):
terrain_height_data = self.game_info.start_raw.terrain_height
bytes_per_pixel = terrain_height_data.bits_per_pixel // 8
width = terrain_height_data.size.x
data = terrain_height_data.data
x, y = position
index = -width * round(y) + round(x)
start = index * bytes_per_pixel
data = data[start: start + bytes_per_pixel]
return int.from_bytes(data, byteorder="little", signed=False)
def geysers(self):
from game_data.units import UnitManager
return UnitManager(self._geysers)
def mineral_fields(self):
from game_data.units import UnitManager
return UnitManager(self._mineral_fields)
def get_available_hubs(self, game_state):
harvesting_hubs = game_state.player_units.filter(build_progress=1, unit_type__in=self.HARVESTING_HUBS)
return UnitManager(list(filter(lambda hub: hub.assigned_harvesters < hub.ideal_harvesters, harvesting_hubs)))
def match_query_output(self, game_state):
town_halls = game_state.player_units.filter(build_progress=1, unit_type__in=self.HARVESTING_HUBS)
return UnitManager(list(filter(lambda th: th.assigned_harvesters > th.ideal_harvesters, town_halls)))