parabellum 0.0.0__py3-none-any.whl → 0.0.73__py3-none-any.whl

parabellum/__init__.py

@@ -1,4 +1,4 @@
-from .env import Parabellum, Scenario
-from .vis import Visualizer
+from . import env, geo, types, utils
+from .env import Env
 
-__all__ = ["Parabellum", "Visualizer", "Scenario"]
+__all__ = ["geo", "env", "types", "utils", "Env"]

parabellum/env.py

@@ -1,296 +1,70 @@
-"""Parabellum environment based on SMAX"""
+# %% env.py
+#   parabellum env
+# by: Noah Syrkis
 
-import jax.numpy as jnp
-import jax
-import numpy as np
-from jax import random
-from jax import jit
-from flax.struct import dataclass
-import chex
-from jaxmarl.environments.smax.smax_env import State, SMAX
-from typing import Tuple, Dict
+# Imports
 from functools import partial
+from typing import Tuple
 
+import jax.numpy as jnp
+import jaxkd as jk
+from jax import random
+from jaxtyping import Array
 
-@dataclass
-class Scenario:
-    """Parabellum scenario"""
-
-    obstacle_coords: chex.Array
-    obstacle_deltas: chex.Array
-
-    unit_types: chex.Array
-    num_allies: int
-    num_enemies: int
-
-    smacv2_position_generation: bool = False
-    smacv2_unit_type_generation: bool = False
-
-
-# default scenario
-scenarios = {
-    "default": Scenario(
-        jnp.array([[6, 10], [26, 10]]) * 8,
-        jnp.array([[0, 12], [0, 1]]) * 8,
-        jnp.zeros((19,), dtype=jnp.uint8),
-        9,
-        10,
-    )
-}
-
-
-class Parabellum(SMAX):
-    def __init__(
-        self,
-        scenario: Scenario = scenarios["default"],
-        unit_type_attack_blasts=jnp.array([0, 0, 0, 0, 0, 0]) + 8,
-        **kwargs,
-    ):
-        super().__init__(scenario=scenario, **kwargs)
-        self.unit_type_attack_blasts = unit_type_attack_blasts
-        self.obstacle_coords = scenario.obstacle_coords.astype(jnp.float32)
-        self.obstacle_deltas = scenario.obstacle_deltas.astype(jnp.float32)
-        self.max_steps = 200
-        # overwrite supers _world_step method
-
-    
-    def _push_units_away(self, state: State, firmness: float = 1.0): # we do it inside the _world_step to allow more obstacles constraints
-        return state
-
-    def _our_push_units_away(self, pos, unit_types, firmness: float = 1.0): # copy of SMAX._push_units_away but used without state and called inside _world_step to allow more obstacles constraints
-        delta_matrix = pos[:, None] - pos[None, :]
-        dist_matrix = (
-            jnp.linalg.norm(delta_matrix, axis=-1)
-            + jnp.identity(self.num_agents)
-            + 1e-6
-        )
-        radius_matrix = (
-            self.unit_type_radiuses[unit_types][:, None]
-            + self.unit_type_radiuses[unit_types][None, :]
-        )
-        overlap_term = jax.nn.relu(radius_matrix / dist_matrix - 1.0)
-        unit_positions = (
-            pos
-            + firmness * jnp.sum(delta_matrix * overlap_term[:, :, None], axis=1) / 2
-        )
-        return unit_positions
-        
-    @partial(jax.jit, static_argnums=(0,))  # replace the _world_step method
-    def _world_step(  # modified version of JaxMARL's SMAX _world_step
-        self,
-        key: chex.PRNGKey,
-        state: State,
-        actions: Tuple[chex.Array, chex.Array],
-    ) -> Tuple[Dict[str, chex.Array], State, Dict[str, float], Dict[str, bool], Dict]:
-        
-        @partial(jax.vmap, in_axes=(None, None, 0, 0))
-        def inter_fn(pos, new_pos, obs, obs_end):
-            d1 = jnp.cross(obs - pos, new_pos - pos)
-            d2 = jnp.cross(obs_end - pos, new_pos - pos)
-            d3 = jnp.cross(pos - obs, obs_end - obs)
-            d4 = jnp.cross(new_pos - obs, obs_end - obs)
-            return (d1 * d2 <= 0) & (d3 * d4 <= 0)
-
-        def update_position(idx, vec):
-            # Compute the movements slightly strangely.
-            # The velocities below are for diagonal directions
-            # because these are easier to encode as actions than the four
-            # diagonal directions. Then rotate the velocity 45
-            # degrees anticlockwise to compute the movement.
-            pos = state.unit_positions[idx]
-            new_pos = (
-                pos
-                + vec
-                * self.unit_type_velocities[state.unit_types[idx]]
-                * self.time_per_step
-            )
-            # avoid going out of bounds
-            new_pos = jnp.maximum(
-                jnp.minimum(new_pos, jnp.array([self.map_width, self.map_height])),
-                jnp.zeros((2,)),
-            )
-
-            #######################################################################
-            ############################################ avoid going into obstacles
-            obs = self.obstacle_coords
-            obs_end = obs + self.obstacle_deltas
-            inters = jnp.any(inter_fn(pos, new_pos, obs, obs_end))
-            new_pos = jnp.where(inters, pos, new_pos)
-
-            #######################################################################
-            #######################################################################
-
-            return new_pos
-
-        #######################################################################
-        ######################################### units close enough to get hit
-
-        def bystander_fn(attacked_idx):
-            idxs = jnp.zeros((self.num_agents,))
-            idxs *= (
-                jnp.linalg.norm(
-                    state.unit_positions - state.unit_positions[attacked_idx], axis=-1
-                )
-                < self.unit_type_attack_blasts[state.unit_types[attacked_idx]]
-            )
-            return idxs
-
-            #######################################################################
-            #######################################################################
-
-        def update_agent_health(idx, action, key):  # TODO: add attack blasts
-            # for team 1, their attack actions are labelled in
-            # reverse order because that is the order they are
-            # observed in
-            attacked_idx = jax.lax.cond(
-                idx < self.num_allies,
-                lambda: action + self.num_allies - self.num_movement_actions,
-                lambda: self.num_allies - 1 - (action - self.num_movement_actions),
-            )
-            # deal with no-op attack actions (i.e. agents that are moving instead)
-            attacked_idx = jax.lax.select(
-                action < self.num_movement_actions, idx, attacked_idx
-            )
-
-            attack_valid = (
-                (
-                    jnp.linalg.norm(
-                        state.unit_positions[idx] - state.unit_positions[attacked_idx]
-                    )
-                    < self.unit_type_attack_ranges[state.unit_types[idx]]
-                )
-                & state.unit_alive[idx]
-                & state.unit_alive[attacked_idx]
-            )
-            attack_valid = attack_valid & (idx != attacked_idx)
-            attack_valid = attack_valid & (state.unit_weapon_cooldowns[idx] <= 0.0)
-            health_diff = jax.lax.select(
-                attack_valid,
-                -self.unit_type_attacks[state.unit_types[idx]],
-                0.0,
-            )
-            # design choice based on the pysc2 randomness details.
-            # See https://github.com/deepmind/pysc2/blob/master/docs/environment.md#determinism-and-randomness
-
-            #########################################################
-            ############################### Add bystander health diff
+from parabellum.types import Action, Config, Obs, State
 
-            bystander_idxs = bystander_fn(attacked_idx)  # TODO: use
-            bystander_valid = (
-                jnp.where(attack_valid, bystander_idxs, jnp.zeros((self.num_agents,)))
-                .astype(jnp.bool_)
-                .astype(jnp.float32)
-            )
-            bystander_health_diff = (
-                bystander_valid * -self.unit_type_attacks[state.unit_types[idx]]
-            )
 
-            #########################################################
-            #########################################################
+# %% Dataclass
+class Env:
+    def init(self, cfg: Config, rng: Array) -> Tuple[Obs, State]:
+        state = init_fn(cfg, rng)  # without jit this takes forever
+        return obs_fn(cfg, state), state
 
-            cooldown_deviation = jax.random.uniform(
-                key, minval=-self.time_per_step, maxval=2 * self.time_per_step
-            )
-            cooldown = (
-                self.unit_type_weapon_cooldowns[state.unit_types[idx]]
-                + cooldown_deviation
-            )
-            cooldown_diff = jax.lax.select(
-                attack_valid,
-                # subtract the current cooldown because we are
-                # going to add it back. This way we effectively
-                # set the new cooldown to `cooldown`
-                cooldown - state.unit_weapon_cooldowns[idx],
-                -self.time_per_step,
-            )
-            return (
-                health_diff,
-                attacked_idx,
-                cooldown_diff,
-                (bystander_health_diff, bystander_idxs),
-            )
+    def step(self, cfg: Config, rng: Array, state: State, action: Action) -> Tuple[Obs, State]:
+        state = step_fn(cfg, rng, state, action)
+        return obs_fn(cfg, state), state
 
-        def perform_agent_action(idx, action, key):
-            movement_action, attack_action = action
-            new_pos = update_position(idx, movement_action)
-            health_diff, attacked_idxes, cooldown_diff, (bystander) = (
-                update_agent_health(idx, attack_action, key)
-            )
 
-            return new_pos, (health_diff, attacked_idxes), cooldown_diff, bystander
+# %% Functions
+def init_fn(cfg: Config, rng: Array) -> State:
+    prob = jnp.ones((cfg.size, cfg.size)).at[cfg.map].set(0).flatten()  # Set
+    flat = random.choice(rng, jnp.arange(prob.size), shape=(cfg.length,), p=prob, replace=True)
+    idxs = (flat // len(cfg.map), flat % len(cfg.map))
+    pos = jnp.float32(jnp.column_stack(idxs))
+    return State(pos=pos, hp=cfg.hp[cfg.types])
 
-        keys = jax.random.split(key, num=self.num_agents)
-        pos, (health_diff, attacked_idxes), cooldown_diff, bystander = jax.vmap(
-            perform_agent_action
-        )(jnp.arange(self.num_agents), actions, keys)
 
-        # checked that no unit passed through an obstacles 
-        new_pos = self._our_push_units_away(pos, state.unit_types)
+def obs_fn(cfg: Config, state: State) -> Obs:  # return info about neighbors ---
+    idxs, dist = jk.extras.query_neighbors_pairwise(state.pos, state.pos, k=cfg.knn)
+    mask = dist < cfg.sight[cfg.types[idxs][:, 0]][..., None]  # | (state.hp[idxs] > 0)
+    pos = (state.pos[idxs] - state.pos[:, None, ...]).at[:, 0, :].set(state.pos) * mask[..., None]
+    args = state.hp, cfg.types, cfg.teams, cfg.reach, cfg.sight, cfg.speed
+    hp, type, team, reach, sight, speed = map(lambda x: x[idxs] * mask, args)
+    return Obs(pos=pos, dist=dist, hp=hp, type=type, team=team, reach=reach, sight=sight, speed=speed, mask=mask)
 
-        # avoid going into obstacles after being pushed 
-        obs = self.obstacle_coords
-        obs_end = obs + self.obstacle_deltas
-   
-        def check_obstacles(pos, new_pos, obs, obs_end):
-            inters = jnp.any(inter_fn(pos, new_pos, obs, obs_end))
-            return jnp.where(inters, pos, new_pos)
-        
-        pos = jax.vmap(check_obstacles, in_axes=(0,0,None,None))(pos, new_pos, obs, obs_end)
-        
-        # Multiple enemies can attack the same unit.
-        # We have `(health_diff, attacked_idx)` pairs.
-        # `jax.lax.scatter_add` aggregates these exactly
-        # in the way we want -- duplicate idxes will have their
-        # health differences added together. However, it is a
-        # super thin wrapper around the XLA scatter operation,
-        # which has this bonkers syntax and requires this dnums
-        # parameter. The usage here was inferred from a test:
-        # https://github.com/google/jax/blob/main/tests/lax_test.py#L2296
-        dnums = jax.lax.ScatterDimensionNumbers(
-            update_window_dims=(),
-            inserted_window_dims=(0,),
-            scatter_dims_to_operand_dims=(0,),
-        )
-        unit_health = jnp.maximum(
-            jax.lax.scatter_add(
-                state.unit_health,
-                jnp.expand_dims(attacked_idxes, 1),
-                health_diff,
-                dnums,
-            ),
-            0.0,
-        )
 
-        #########################################################
-        ############################ subtracting bystander health
+def step_fn(cfg: Config, rng: Array, state: State, action: Action) -> State:
+    idx, norm = jk.extras.query_neighbors_pairwise(state.pos + action.pos, state.pos, k=2)
+    args = rng, cfg, state, action, idx, norm
+    return State(pos=partial(push_fn, cfg, rng, idx, norm)(move_fn(*args)), hp=blast_fn(*args))  # type: ignore
 
-        _, bystander_health_diff = bystander
-        unit_health -= bystander_health_diff.sum(axis=0)  # might be axis=1
 
-        #########################################################
-        #########################################################
+def move_fn(rng: Array, cfg: Config, state: State, action: Action, idx: Array, norm: Array) -> Array:
+    speed = cfg.speed[cfg.types][..., None]  # max speed of a unit (step size, really)
+    pos = state.pos + action.pos.clip(-speed, speed) * action.move[..., None]  # new poss
+    mask = ((pos < 0).any(axis=-1) | ((pos >= cfg.size).any(axis=-1)) | (cfg.map[*jnp.int32(pos).T] > 0))[..., None]
+    return jnp.where(mask, state.pos, pos)  # compute new position
 
-        unit_weapon_cooldowns = state.unit_weapon_cooldowns + cooldown_diff
-        state = state.replace(
-            unit_health=unit_health,
-            unit_positions=pos,
-            unit_weapon_cooldowns=unit_weapon_cooldowns,
-        )
-        return state
 
-if __name__ == "__main__":
-    env = Parabellum(map_width=256, map_height=256)
-    rng, key = random.split(random.PRNGKey(0))
-    obs, state = env.reset(key)
-    state_seq = []
-    for step in range(100):
-        rng, key = random.split(rng)
-        key_act = random.split(key, len(env.agents))
-        actions = {
-            agent: jax.random.randint(key_act[i], (), 0, 5)
-            for i, agent in enumerate(env.agents)
-        }
-        _, state, _, _, _ = env.step(key, state, actions)
-        state_seq.append((obs, state, actions))
+def blast_fn(rng: Array, cfg: Config, state: State, action: Action, idx: Array, norm: Array) -> Array:
+    dam = (cfg.dam[cfg.types] * action.cast)[..., None] * jnp.ones_like(idx)
+    return state.hp - jnp.zeros(cfg.length, dtype=jnp.int32).at[idx.flatten()].add(dam.flatten())
 
 
+def push_fn(cfg: Config, rng: Array, idx: Array, norm: Array, pos: Array) -> Array:
+    return pos + random.normal(rng, pos.shape) * 0.1
+    # params need to be tweaked, and matched with unit size
+    pos_diff = pos[:, None, :] - pos[idx]  # direction away from neighbors
+    mask = (norm < cfg.r[cfg.types][..., None]) & (norm > 0)
+    pos = pos + jnp.where(mask[..., None], pos_diff * cfg.force / (norm[..., None] + 1e-6), 0.0).sum(axis=1)
+    return pos + random.normal(rng, pos.shape) * 0.1

parabellum/geo.py

@@ -0,0 +1,213 @@
+# %% geo.py
+#   script for geospatial level generation
+# by: Noah Syrkis
+
+# %% Imports
+from rasterio import features, transform
+
+# from jax import tree
+from geopy.geocoders import Nominatim
+from geopy.distance import distance
+import contextily as cx
+import jax.numpy as jnp
+import cartopy.crs as ccrs
+from jaxtyping import Array
+from shapely import box
+import osmnx as ox
+import geopandas as gpd
+from collections import namedtuple
+from typing import Tuple
+import matplotlib.pyplot as plt
+from cachier import cachier
+# from jax.scipy.signal import convolve
+# from parabellum.types import Terrain
+
+# %% Types
+Coords = Tuple[float, float]
+BBox = namedtuple("BBox", ["north", "south", "east", "west"])  # type: ignore
+
+# %% Constants
+provider = cx.providers.Stadia.StamenTerrain(  # type: ignore
+    api_key="86d0d32b-d2fe-49af-8db8-f7751f58e83f"
+)
+provider["url"] = provider["url"] + "?api_key={api_key}"
+tags = {
+    "building": True,
+    "water": True,
+    "highway": True,
+    "landuse": [
+        "grass",
+        "forest",
+        "flowerbed",
+        "greenfield",
+        "village_green",
+        "recreation_ground",
+    ],
+    "leisure": "garden",
+}  #  "road": True}
+
+
+# %% Coordinate function
+def get_coordinates(place: str) -> Coords:
+    geolocator = Nominatim(user_agent="parabellum")
+    point = geolocator.geocode(place)
+    return point.latitude, point.longitude  # type: ignore
+
+
+def get_bbox(place: str, buffer) -> BBox:
+    """Get bounding box from place name in crs 4326."""
+    coords = get_coordinates(place)
+    north = distance(meters=buffer).destination(coords, bearing=0).latitude
+    south = distance(meters=buffer).destination(coords, bearing=180).latitude
+    east = distance(meters=buffer).destination(coords, bearing=90).longitude
+    west = distance(meters=buffer).destination(coords, bearing=270).longitude
+    return BBox(north, south, east, west)  # type: ignore
+
+
+def basemap_fn(bbox: BBox, gdf) -> Array:
+    fig, ax = plt.subplots(figsize=(20, 20), subplot_kw={"projection": ccrs.Mercator()})
+    gdf.plot(ax=ax, color="black", alpha=0, edgecolor="black")  # type: ignore
+    cx.add_basemap(ax, crs=gdf.crs, source=provider, zoom="auto")  # type: ignore
+    bbox = gdf.total_bounds
+    ax.set_extent([bbox[0], bbox[2], bbox[1], bbox[3]], crs=ccrs.Mercator())  # type: ignore
+    plt.axis("off")
+    plt.tight_layout(pad=0)
+    fig.canvas.draw()
+    image = jnp.array(fig.canvas.renderer._renderer)  # type: ignore
+    plt.close(fig)
+    return image
+
+
+@cachier()
+def geography_fn(place, buffer):
+    bbox = get_bbox(place, buffer)
+    map_data = ox.features_from_bbox(bbox=bbox, tags=tags)
+    gdf = gpd.GeoDataFrame(map_data)
+    gdf = gdf.clip(box(bbox.west, bbox.south, bbox.east, bbox.north)).to_crs("EPSG:3857")
+    raster = raster_fn(gdf, shape=(buffer, buffer))
+    # basemap = jnp.rot90(basemap_fn(bbox, gdf), 3)
+    # kernel = jnp.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]])
+    trans = lambda x: jnp.bool(x)  # jnp.rot90(x, 3)  # noqa
+    terrain = trans(raster[0])  # Terrain(
+    #        building=trans(raster[0]),
+    #        water=trans(raster[1] - convolve(raster[1] * raster[2], kernel, mode="same") > 0),
+    #        forest=trans(jnp.logical_or(raster[3], raster[4])),
+    #        basemap=basemap,
+    #    )
+    # terrain = tree.map(lambda x: x.astype(jnp.int16), terrain)
+    return terrain
+
+
+# =======
+#     terrain = tps.Terrain(building=trans(raster[0] - convolve(raster[0]*raster[2], kernel, mode='same')>0),
+#                           water=trans(raster[1] - convolve(raster[1]*raster[2], kernel, mode='same')>0),
+#                           forest=trans(jnp.logical_or(raster[3], raster[4])),
+#                           basemap=basemap)
+#     return terrain, gdf
+# >>>>>>> aeb13033e57083cc512a60f8f60a3db47a65ac32
+
+
+def raster_fn(gdf, shape) -> Array:
+    bbox = gdf.total_bounds
+    t = transform.from_bounds(*bbox, *shape)  # type: ignore
+    raster = jnp.array([feature_fn(t, feature, gdf, shape) for feature in tags])
+    return raster
+
+
+def feature_fn(t, feature, gdf, shape):
+    if feature not in gdf.columns:
+        return jnp.zeros(shape)
+    gdf = gdf[~gdf[feature].isna()]
+    raster = features.rasterize(gdf.geometry, out_shape=shape, transform=t, fill=0)  # type: ignore
+    return raster
+
+
+# %%
+# def normalize(x):
+#     return (np.array(x) - m) / (M - m)
+
+
+# def get_bridges(gdf):
+#     xmin, ymin, xmax, ymax = gdf.total_bounds
+#     m = np.array([xmin, ymin])
+#     M = np.array([xmax, ymax])
+
+#     bridges = {}
+#     for idx, bridge in gdf[gdf["bridge"] == "yes"].iterrows():
+#         if type(bridge["name"]) == str:
+#             bridges[idx[1]] = {
+#                 "name": bridge["name"],
+#                 "coords": normalize(
+#                     [bridge.geometry.centroid.x, bridge.geometry.centroid.y]
+#                 ),
+#             }
+#     return bridges
+
+
+"""
+# %%
+if __name__ == "__main__":
+    place = "Thun, Switzerland"
+<<<<<<< HEAD
+    terrain = geography_fn(place, 300)
+
+=======
+    terrain, gdf = geography_fn(place, 300)
+
+>>>>>>> aeb13033e57083cc512a60f8f60a3db47a65ac32
+    fig, axes = plt.subplots(1, 5, figsize=(20, 20))
+    axes[0].imshow(jnp.rot90(terrain.building), cmap="gray")
+    axes[1].imshow(jnp.rot90(terrain.water), cmap="gray")
+    axes[2].imshow(jnp.rot90(terrain.forest), cmap="gray")
+    axes[3].imshow(jnp.rot90(terrain.building + terrain.water + terrain.forest))
+    axes[4].imshow(jnp.rot90(terrain.basemap))
+
+    # %%
+    W, H, _  = terrain.basemap.shape
+    bridges = get_bridges(gdf)
+
+    # %%
+    print("Bridges:")
+    for bridge in bridges.values():
+        x, y = int(bridge["coords"][0]*300), int(bridge["coords"][1]*300)
+        print(bridge["name"], f"at ({x}, {y})")
+
+    # %%
+    plt.subplots(figsize=(7,7))
+    plt.imshow(jnp.rot90(terrain.basemap))
+    X = [b["coords"][0]*W for b in bridges.values()]
+    Y = [(1-b["coords"][1])*H for b in bridges.values()]
+    plt.scatter(X, Y)
+    for i in range(len(X)):
+        x,y = int(X[i]), int(Y[i])
+        plt.text(x, y, str((int(x/W*300), int((1-(y/H))*300))))
+
+# %%
+
+# %% [raw]
+# fig, ax = plt.subplots(figsize=(10, 10))
+# gdf.plot(ax=ax, color='lightgray')  # Plot all features
+# bridges.plot(ax=ax, color='red')     # Highlight bridges in red
+# plt.show()
+
+# %%
+
+"""
+
+# BBox = namedtuple("BBox", ["north", "south", "east", "west"])  # type: ignore
+
+
+# def to_mercator(bbox: BBox) -> BBox:
+# proj = ccrs.Mercator()
+# west, south = proj.transform_point(bbox.west, bbox.south, ccrs.PlateCarree())
+# east, north = proj.transform_point(bbox.east, bbox.north, ccrs.PlateCarree())
+# return BBox(north=north, south=south, east=east, west=west)
+#
+#
+# def to_platecarree(bbox: BBox) -> BBox:
+# proj = ccrs.PlateCarree()
+# west, south = proj.transform_point(bbox.west, bbox.south, ccrs.Mercator())
+#
+# east, north = proj.transform_point(bbox.east, bbox.north, ccrs.Mercator())
+# return BBox(north=north, south=south, east=east, west=west)
+#