parabellum 0.2.23__py3-none-any.whl → 0.2.25__py3-none-any.whl

parabellum/env.py

@@ -19,6 +19,7 @@ class Scenario:
 
     place: str
     terrain_raster: jnp.ndarray
+    unit_starting_sectors: jnp.ndarray
     unit_types: chex.Array
     num_allies: int
     num_enemies: int
@@ -45,6 +46,7 @@ scenarios = {
     "default": Scenario(
         "Identity Town",
         jnp.eye(64, dtype=jnp.uint8),
+        jnp.array([[0, 0, 0.2, 0.2], [0.7,0.7,0.2,0.2]]),
         jnp.zeros((19,), dtype=jnp.uint8),
         9,
         10,
@@ -52,11 +54,20 @@ scenarios = {
 }
 
 
-def make_scenario(place, terrain_raster, num_allies=9, num_enemies=10):
-    """Create a scenario"""
-    num_agents = num_allies + num_enemies
-    unit_types = jnp.zeros((num_agents,)).astype(jnp.uint8)
-    return Scenario(place, terrain_raster, unit_types, num_allies, num_enemies)
+def make_scenario(place, terrain_raster, unit_starting_sectors, allies_type, n_allies, enemies_type, n_enemies):
+    if type(allies_type) == int:
+        allies = [allies_type] * n_allies
+    else:
+        assert(len(allies_type) == n_allies)
+        allies = allies_type
+        
+    if type(enemies_type) == int:
+        enemies = [enemies_type] * n_enemies
+    else:
+        assert(len(enemies_type) == n_enemies)
+        enemies = enemies_type
+    unit_types = jnp.array(allies + enemies, dtype=jnp.uint8)
+    return Scenario(place, terrain_raster, unit_starting_sectors, unit_types, n_allies, n_enemies)
 
 
 def spawn_fn(pool, offset: jnp.ndarray, n: int, rng: jnp.ndarray):
@@ -81,29 +92,41 @@ def sector_fn(terrain: jnp.ndarray, sector_id: int):
     return jnp.nonzero(sector), offset
 
 
+def sector_fn(terrain: jnp.ndarray, sector: jnp.ndarray):
+    """return sector slice of terrain"""
+    width, height = terrain.shape
+    coordx, coordy = int(sector[0] * width), int(sector[1] * height)
+    sector = terrain[coordy : coordy + int(sector[3] * height), coordx : coordx + int(sector[2] * width)] == 0
+    offset = jnp.array([coordx, coordy])
+    # sector is jnp.nonzero
+    return jnp.nonzero(sector.T), offset
+
+
 class Environment(SMAX):
     def __init__(self, scenario: Scenario, **kwargs):
         map_height, map_width = scenario.terrain_raster.shape
         args = dict(scenario=scenario, map_height=map_height, map_width=map_width)
         super(Environment, self).__init__(**args, walls_cause_death=False, **kwargs)
         self.terrain_raster = scenario.terrain_raster
+        self.unit_starting_sectors = scenario.unit_starting_sectors
         # self.unit_type_names = ["tinker", "tailor", "soldier", "spy"]
         # self.unit_type_health = jnp.array([100, 100, 100, 100], dtype=jnp.float32)
         # self.unit_type_damage = jnp.array([10, 10, 10, 10], dtype=jnp.float32)
         self.scenario = scenario
+        self.unit_type_velocities=jnp.array([3.15, 2.25, 4.13, 3.15, 4.13, 3.15])/2.5
         self.unit_type_attack_blasts = jnp.zeros((3,), dtype=jnp.float32)  # TODO: add
         self.max_steps = 200
         self._push_units_away = lambda state, firmness = 1: state  # overwrite push units
-        self.top_sector, self.top_sector_offset = sector_fn(self.terrain_raster, 0)
-        self.low_sector, self.low_sector_offset = sector_fn(self.terrain_raster, 24)
+        self.team0_sector, self.team0_sector_offset = sector_fn(self.terrain_raster, self.unit_starting_sectors[0]) # sector_fn(self.terrain_raster, 0)
+        self.team1_sector, self.team1_sector_offset = sector_fn(self.terrain_raster, self.unit_starting_sectors[1]) # sector_fn(self.terrain_raster, 24)
 
 
     @partial(jax.jit, static_argnums=(0,))
     def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:
         """Environment-specific reset."""
         ally_key, enemy_key = jax.random.split(rng)
-        team_0_start = spawn_fn(self.top_sector, self.top_sector_offset, self.num_allies, ally_key)
-        team_1_start = spawn_fn(self.low_sector, self.low_sector_offset, self.num_enemies, enemy_key)
+        team_0_start = spawn_fn(self.team0_sector, self.team0_sector_offset, self.num_allies, ally_key)
+        team_1_start = spawn_fn(self.team1_sector, self.team1_sector_offset, self.num_enemies, enemy_key)
         unit_positions = jnp.concatenate([team_0_start, team_1_start])
         unit_teams = jnp.zeros((self.num_agents,))
         unit_teams = unit_teams.at[self.num_allies :].set(1)
@@ -126,14 +149,9 @@ class Environment(SMAX):
         state = self._push_units_away(state)  # type: ignore
         obs = self.get_obs(state)
         world_state = self.get_world_state(state)
-        # obs["world_state"] = jax.lax.stop_gradient(world_state)
+        obs["world_state"] = jax.lax.stop_gradient(world_state)
         return obs, state
 
-    def step_env(self, rng, state: State, action: Array):
-        obs, state, rewards, dones, infos = super().step_env(rng, state, action)
-        # delete world_state from obs
-        obs.pop("world_state")
-        return obs, state, rewards, dones, infos
 
     def _our_push_units_away(
         self, pos, unit_types, firmness: float = 1.0
@@ -172,15 +190,15 @@ class Environment(SMAX):
 
         def raster_crossing(pos, new_pos):
             pos, new_pos = pos.astype(jnp.int32), new_pos.astype(jnp.int32)
-            raster = self.terrain_raster
-            axis = jnp.argmax(jnp.abs(new_pos - pos), axis=-1)
-            minimum = jnp.minimum(pos[axis], new_pos[axis]).squeeze()
-            maximum = jnp.maximum(pos[axis], new_pos[axis]).squeeze()
-            segment = jnp.where(axis == 0, raster[pos[1]], raster.T[pos[0]])
-            segment = jnp.where(jnp.arange(segment.shape[0]) >= minimum, segment, 0)
-            segment = jnp.where(jnp.arange(segment.shape[0]) <= maximum, segment, 0)
-            return jnp.any(segment)
-
+            raster = jnp.copy(self.terrain_raster)
+            minimum = jnp.minimum(pos, new_pos)
+            maximum = jnp.maximum(pos, new_pos)
+            raster = jnp.where(jnp.arange(raster.shape[0]) >= minimum[0], raster, 0)
+            raster = jnp.where(jnp.arange(raster.shape[0]) <= maximum[0], raster, 0)
+            raster = jnp.where(jnp.arange(raster.shape[1]) >= minimum[1], raster.T, 0).T
+            raster = jnp.where(jnp.arange(raster.shape[1]) <= maximum[1], raster.T, 0).T
+            return jnp.any(raster)
+        
         def update_position(idx, vec):
             # Compute the movements slightly strangely.
             # The velocities below are for diagonal directions
@@ -202,11 +220,7 @@ class Environment(SMAX):
 
             #######################################################################
             ############################################ avoid going into obstacles
-            """ obs = self.obstacle_coords
-            obs_end = obs + self.obstacle_deltas
-            inters = jnp.any(intersect_fn(pos, new_pos, obs, obs_end)) """
             clash = raster_crossing(pos, new_pos)
-            # flag = jnp.logical_or(inters, rastersects)
             new_pos = jnp.where(clash, pos, new_pos)
 
             #######################################################################
@@ -319,39 +333,11 @@ class Environment(SMAX):
 
         # units push each other
         new_pos = self._our_push_units_away(pos, state.unit_types)
-
-        # avoid going into obstacles after being pushed
-
-        bondaries_coords = jnp.array(
-            [[0, 0], [0, 0], [self.map_width, 0], [0, self.map_height]]
-        )
-        bondaries_deltas = jnp.array(
-            [
-                [self.map_width, 0],
-                [0, self.map_height],
-                [0, self.map_height],
-                [self.map_width, 0],
-            ]
-        )
-        """ obstacle_coords = jnp.concatenate(
-            [self.obstacle_coords, bondaries_coords]
-        )  # add the map boundaries to the obstacles to avoid
-        obstacle_deltas = jnp.concatenate(
-            [self.obstacle_deltas, bondaries_deltas]
-        )  # add the map boundaries to the obstacles to avoid
-        obst_start = obstacle_coords
-        obst_end = obst_start + obstacle_deltas """
-
-        def check_obstacles(pos, new_pos, obst_start, obst_end):
-            intersects = jnp.any(intersect_fn(pos, new_pos, obst_start, obst_end))
-            rastersect = raster_crossing(pos, new_pos)
-            flag = jnp.logical_or(intersects, rastersect)
-            return jnp.where(flag, pos, new_pos)
-
-        """ pos = jax.vmap(check_obstacles, in_axes=(0, 0, None, None))(
-            pos, new_pos, obst_start, obst_end
-        ) """
-
+        clash = jax.vmap(raster_crossing)(pos, new_pos)
+        pos = jax.vmap(jnp.where)(clash, pos, new_pos)
+        # avoid going out of bounds
+        pos = jnp.maximum(jnp.minimum(pos, jnp.array([self.map_width, self.map_height])),jnp.zeros((2,)),)
+        
         # Multiple enemies can attack the same unit.
         # We have `(health_diff, attacked_idx)` pairs.
         # `jax.lax.scatter_add` aggregates these exactly

parabellum/map.py

@@ -14,6 +14,8 @@ import rasterio.transform
 from typing import Optional, Tuple
 from geopy.location import Location
 from shapely.geometry import Point
+import os
+import pickle 
 
 # constants
 geolocator = Nominatim(user_agent="parabellum")
@@ -36,16 +38,41 @@ def rasterize_geometry(gdf: gpd.GeoDataFrame, size: int) -> jnp.ndarray:
     w, s, e, n = gdf.total_bounds
     transform = rasterio.transform.from_bounds(w, s, e, n, size, size)
     raster = features.rasterize(gdf.geometry, out_shape=(size, size), transform=transform)
-    return jnp.array(jnp.rot90(raster, 2)).astype(jnp.uint8)
+    return jnp.array(jnp.flip(raster, 0) ).astype(jnp.uint8)
 
-def terrain_fn(place: str, size: int = 1000) -> Tuple[jnp.ndarray, jnp.ndarray]:
+# +
+def get_from_cache(place, size):
+    if os.path.exists("./cache"):
+        name = str(hash((place, size))) + ".pk"
+        if os.path.exists("./cache/" + name):
+            with open("./cache/" + name, "rb") as f:
+                (mask, base) = pickle.load(f)
+            return (mask, base.astype(jnp.int64))
+    return (None, None)
+    
+def save_in_cache(place, size, mask, base):
+    if not os.path.exists("./cache"):
+        os.makedirs("./cache")
+    name = str(hash((place, size))) + ".pk"
+    with open("./cache/" + name, "wb") as f:
+        pickle.dump((mask, base), f)
+
+def terrain_fn(place: str, size: int = 1000, with_cache: bool = True) -> Tuple[jnp.ndarray, jnp.ndarray]:
     """Returns a rasterized map of buildings for a given location."""
-    point = get_location(place)
-    gdf = get_building_geometry(point, size)
-    mask = rasterize_geometry(gdf, size)
-    base = get_basemap(place, size)
+    if with_cache:
+        mask, base = get_from_cache(place, size)
+    if not with_cache or mask is None:
+        point = get_location(place)
+        gdf = get_building_geometry(point, size)
+        mask = rasterize_geometry(gdf, size)
+        base = get_basemap(place, size)
+        if with_cache:
+            save_in_cache(place, size, mask, base)
     return mask, base
 
+
+# -
+
 def get_basemap(place: str, size: int = 1000) -> jnp.ndarray:
     """Returns a basemap for a given place as a JAX array."""
     point = get_location(place)
@@ -54,15 +81,20 @@ def get_basemap(place: str, size: int = 1000) -> jnp.ndarray:
     # get the middle size x size square
     basemap = basemap[(basemap.shape[0] - size) // 2:(basemap.shape[0] + size) // 2,
                         (basemap.shape[1] - size) // 2:(basemap.shape[1] + size) // 2]
-    return jnp.array(jnp.rot90(basemap, 2)).astype(jnp.uint8)
+    return basemap # jnp.array(jnp.rot90(basemap, 2)).astype(jnp.uint8)
 
 
 if __name__ == "__main__":
-    import seaborn as sns
-    place = "Thun, Switzerland"
-    mask, base = terrain_fn(place)
-
-    fig, ax = plt.subplots(1, 2, figsize=(10, 5))
-    ax[0].imshow(mask) # type: ignore
+    place = "Cauvicourt, 14190, France"
+    mask, base = terrain_fn(place, 500)
+    
+    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
+    ax[0].imshow(jnp.flip(mask,0)) # type: ignore
     ax[1].imshow(base) # type: ignore
+    ax[2].imshow(base) # type: ignore
+    ax[2].imshow(jnp.flip(mask,0), alpha=jnp.flip(mask,0)) # type: ignore
     plt.show()
+
+
+
+

parabellum/vis.py

@@ -56,7 +56,6 @@ class Visualizer(SMAXVisualizer):
         self.skin.scale = self.skin.size / env.map_width  # assumes square map
         self.env = env
 
-
     def animate(self, save_fname: Optional[str] = "output/parabellum", view=None):
         expanded_state_seq, expanded_action_seq = expand_fn(self.env, self.state_seq, self.action_seq)
         state_seq_seq, action_seq_seq = unbatch_fn(expanded_state_seq, expanded_action_seq)
@@ -86,7 +85,8 @@ def init_frame(env, skin, image, state: pb.State, action: Array, idx: int) -> py
 
 
 def transform_frame(env, skin, frame):
-    frame = np.rot90(pygame.surfarray.pixels3d(frame).swapaxes(0, 1), 2)
+    #frame = np.rot90(pygame.surfarray.pixels3d(frame).swapaxes(0, 1), 2)
+    frame = np.flip(pygame.surfarray.pixels3d(frame).swapaxes(0, 1), 0)
     return frame
 
 
@@ -136,7 +136,7 @@ def text_fn(text):
 
 def image_fn(skin: Skin):  # TODO:
     """Create an image for background (basemap or maskmap)"""
-    motif = cv2.resize(np.array(skin.maskmap.T), (skin.size, skin.size), interpolation=cv2.INTER_LANCZOS4).astype(np.uint8)
+    motif = cv2.resize(np.array(skin.maskmap.T), (skin.size, skin.size), interpolation=cv2.INTER_NEAREST).astype(np.uint8)
     motif = (motif > 0).astype(np.uint8)
     image = np.zeros((skin.size, skin.size, 3), dtype=np.uint8) + skin.bg
     image[motif == 1] = skin.fg

{parabellum-0.2.23.dist-info → parabellum-0.2.25.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: parabellum
-Version: 0.2.23
+Version: 0.2.25
 Summary: Parabellum environment for parallel warfare simulation
 Home-page: https://github.com/syrkis/parabellum
 License: MIT

parabellum-0.2.25.dist-info/RECORD

@@ -0,0 +1,10 @@
+parabellum/__init__.py,sha256=-5cWXJkHnfH_CbhTEall8Wak8McAFXZHP1L8Fu7Uo5k,373
+parabellum/aid.py,sha256=HWST27inTFXcp8b11izJF0U7N7DZnRTIS3n1Qfa-Ko4,106
+parabellum/env.py,sha256=H1VwxHj8KqbjBJ8b7NMxkAg3Q0qwVzGulrigc26Tzkc,16663
+parabellum/gun.py,sha256=nvsJdcZ2Qd6lbPlAgsUiaLhstTi1UdLQ8kOnbCenucY,2618
+parabellum/map.py,sha256=UwMqwySasX5oLw9v5YMJARAPwvQThLTRW36NpbwvBC8,3564
+parabellum/run.py,sha256=EO_F7VPwayatpSHrcbSahtinsV4QObhcx0jo-4KZO1E,3472
+parabellum/vis.py,sha256=q7_OIMjzt-7nBOojVVW7Wiiq9ojsjaltIsH6eOOxPKk,6116
+parabellum-0.2.25.dist-info/METADATA,sha256=UCuRLYhSUxnebs5pRs1FaWK67PUJw0kgsJ2UaaxNM9Q,2671
+parabellum-0.2.25.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.2.25.dist-info/RECORD,,

parabellum-0.2.23.dist-info/RECORD

@@ -1,10 +0,0 @@
-parabellum/__init__.py,sha256=-5cWXJkHnfH_CbhTEall8Wak8McAFXZHP1L8Fu7Uo5k,373
-parabellum/aid.py,sha256=HWST27inTFXcp8b11izJF0U7N7DZnRTIS3n1Qfa-Ko4,106
-parabellum/env.py,sha256=2IuXipHXUJAyfrjPtDW7uINb8mY4G3-xQ_lhReMGqLs,16985
-parabellum/gun.py,sha256=nvsJdcZ2Qd6lbPlAgsUiaLhstTi1UdLQ8kOnbCenucY,2618
-parabellum/map.py,sha256=EUcPe4Upu9MQzS8h15IVPGCaAyRPLSkmoLd5ZT-V4Pk,2599
-parabellum/run.py,sha256=EO_F7VPwayatpSHrcbSahtinsV4QObhcx0jo-4KZO1E,3472
-parabellum/vis.py,sha256=uXTnhJL23JLQHW9by-M4bF73dSVA5TIkpNdfo_Go2Ro,6045
-parabellum-0.2.23.dist-info/METADATA,sha256=qwEcFJksQ54_MNpMj8Zq1EC2qDgltEC9u9ivZFxZ0Eg,2671
-parabellum-0.2.23.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-parabellum-0.2.23.dist-info/RECORD,,