diff --git a/urnai/sc2/states/__init__.py b/urnai/sc2/states/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/urnai/sc2/states/terran_state.py b/urnai/sc2/states/terran_state.py new file mode 100644 index 00000000..c3bd47a5 --- /dev/null +++ b/urnai/sc2/states/terran_state.py @@ -0,0 +1,121 @@ +import math + +import numpy as np +from pysc2.env import sc2_env +from pysc2.lib import features, units + +from urnai.states.state_base import StateBase + + +class TerranState(StateBase): + + def __init__(self, grid_size=4): + + self.grid_size = grid_size + # Size of the state returned + # 22: number of data added to the state (number of minerals, army_count, etc) + # 2 * 4 ** 2: 2 grids of size 4 x 4, representing enemy and player units + self._state_size = int(22 + 2 * (self.grid_size ** 2)) + self.player_race = sc2_env.Race.terran + self.base_top_left = None + self._state = None + + def update(self, obs): + if obs.game_loop[0] < 80 and self.base_top_left is None: + + commandcenter = get_units_by_type(obs, units.Terran.CommandCenter) + + if len(commandcenter) > 0: + townhall = commandcenter[0] + self.player_race = sc2_env.Race.terran + + self.base_top_left = (townhall.x < 32) + + new_state = [] + # Adds general information from the player. + new_state.append(obs.player.minerals / 6000) + new_state.append(obs.player.vespene / 6000) + new_state.append(obs.player.food_cap / 200) + new_state.append(obs.player.food_used / 200) + new_state.append(obs.player.food_army / 200) + new_state.append(obs.player.food_workers / 200) + new_state.append((obs.player.food_cap - obs.player.food_used) / 200) + new_state.append(obs.player.army_count / 200) + new_state.append(obs.player.idle_worker_count / 200) + + # Adds information related to player's Terran units/buildings. + new_state.append(get_my_units_amount(obs, units.Terran.CommandCenter) + + get_my_units_amount(obs, units.Terran.OrbitalCommand) + + get_my_units_amount(obs, units.Terran.PlanetaryFortress) / 2) + new_state.append(get_my_units_amount(obs, units.Terran.SupplyDepot) / 18) + new_state.append(get_my_units_amount(obs, units.Terran.Refinery) / 4) + new_state.append(get_my_units_amount(obs, units.Terran.EngineeringBay)) + new_state.append(get_my_units_amount(obs, units.Terran.Armory)) + new_state.append(get_my_units_amount(obs, units.Terran.MissileTurret) / 4) + new_state.append(get_my_units_amount(obs, units.Terran.SensorTower)/1) + new_state.append(get_my_units_amount(obs, units.Terran.Bunker)/4) + new_state.append(get_my_units_amount(obs, units.Terran.FusionCore)) + new_state.append(get_my_units_amount(obs, units.Terran.GhostAcademy)) + new_state.append(get_my_units_amount(obs, units.Terran.Barracks) / 3) + new_state.append(get_my_units_amount(obs, units.Terran.Factory) / 2) + new_state.append(get_my_units_amount(obs, units.Terran.Starport) / 2) + + # Instead of making a vector for all coordnates on the map, we'll + # discretize our enemy space + # and use a 4x4 grid to store enemy positions by marking a square as 1 if + # there's any enemy on it. + + enemy_grid = np.zeros((self.grid_size, self.grid_size)) + player_grid = np.zeros((self.grid_size, self.grid_size)) + + enemy_units = [unit for unit in obs.raw_units if + unit.alliance == features.PlayerRelative.ENEMY] + player_units = [unit for unit in obs.raw_units if + unit.alliance == features.PlayerRelative.SELF] + + for i in range(0, len(enemy_units)): + y = int(math.ceil((enemy_units[i].x + 1) / 64 / self.grid_size)) + x = int(math.ceil((enemy_units[i].y + 1) / 64 / self.grid_size)) + enemy_grid[x - 1][y - 1] += 1 + + for i in range(0, len(player_units)): + y = int(math.ceil((player_units[i].x + 1) / (64 / self.grid_size))) + x = int(math.ceil((player_units[i].y + 1) / (64 / self.grid_size))) + player_grid[x - 1][y - 1] += 1 + + if not self.base_top_left: + enemy_grid = np.rot90(enemy_grid, 2) + player_grid = np.rot90(player_grid, 2) + + # Normalizing the values to always be between 0 and 1 + # (since the max amount of units in SC2 is 200) + enemy_grid = enemy_grid / 200 + player_grid = player_grid / 200 + + new_state.extend(enemy_grid.flatten()) + new_state.extend(player_grid.flatten()) + final_state = np.expand_dims(new_state, axis=0) + + self._state = final_state + return final_state + + @property + def state(self): + return self._state + + @property + def dimension(self): + return self._state_size + + def reset(self): + self._state = None + self.base_top_left = None + +def get_my_units_amount(obs, unit_type): + return len(get_units_by_type(obs, unit_type, features.PlayerRelative.SELF)) + +def get_units_by_type(obs, unit_type, alliance=features.PlayerRelative.SELF): + return [unit for unit in obs.raw_units + if unit.unit_type == unit_type + and unit.alliance == alliance + and unit.build_progress == 100] \ No newline at end of file