parabellum 0.2.20__py3-none-any.whl → 0.2.21__py3-none-any.whl

parabellum/__init__.py

@@ -1,12 +1,18 @@
-from .env import Environment, Scenario, scenarios, make_scenario
-from .vis import Visualizer
+from .env import Environment, Scenario, scenarios, make_scenario, State
+from .vis import Visualizer, Skin
 from .map import terrain_fn
+from .gun import bullet_fn
+# from .aid import aid
+# from .run import run
 
 __all__ = [
     "Environment",
     "Scenario",
     "scenarios",
     "make_scenario",
+    "State",
     "Visualizer",
+    "Skin",
     "terrain_fn",
+    "bullet_fn",
 ]

parabellum/aid.py

@@ -0,0 +1,5 @@
+# aid.py
+#   what you call utils.py when you want file names to be 3 letters
+# by: Noah Syrkis
+
+# imports

parabellum/env.py

@@ -3,13 +3,13 @@
 import jax.numpy as jnp
 import jax
 import numpy as np
-from jax import random
+from jax import random, Array
 from jax import jit
 from flax.struct import dataclass
 import chex
 from jax import vmap
-from jaxmarl.environments.smax.smax_env import State, SMAX
-from typing import Tuple, Dict
+from jaxmarl.environments.smax.smax_env import SMAX
+from typing import Tuple, Dict, cast
 from functools import partial
 
 
@@ -18,7 +18,7 @@ class Scenario:
     """Parabellum scenario"""
 
     place: str
-    terrain_raster: chex.Array
+    terrain_raster: jnp.ndarray
     unit_types: chex.Array
     num_allies: int
     num_enemies: int
@@ -26,6 +26,19 @@ class Scenario:
     smacv2_position_generation: bool = False
     smacv2_unit_type_generation: bool = False
 
+@dataclass
+class State:
+    unit_positions: Array
+    unit_alive: Array
+    unit_teams: Array
+    unit_health: Array
+    unit_types: Array
+    unit_weapon_cooldowns: Array
+    prev_movement_actions: Array
+    prev_attack_actions: Array
+    time: int
+    terminal: bool
+
 
 # default scenario
 scenarios = {
@@ -46,7 +59,7 @@ def make_scenario(place, terrain_raster, num_allies=9, num_enemies=10):
     return Scenario(place, terrain_raster, unit_types, num_allies, num_enemies)
 
 
-def spawn_fn(pool: jnp.ndarray, offset: jnp.ndarray, n: int, rng: jnp.ndarray):
+def spawn_fn(pool, offset: jnp.ndarray, n: int, rng: jnp.ndarray):
     """Spawns n agents on a map."""
     rng, key_start, key_noise = random.split(rng, 3)
     noise = random.uniform(key_noise, (n, 2)) * 0.5
@@ -80,22 +93,22 @@ class Environment(SMAX):
         self.scenario = scenario
         self.unit_type_attack_blasts = jnp.zeros((3,), dtype=jnp.float32)  # TODO: add
         self.max_steps = 200
-        self._push_units_away = lambda x: x  # overwrite push units
-        self.top_sector = sector_fn(self.terrain_raster, 0)
-        self.low_sector = sector_fn(self.terrain_raster, 24)
+        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)
 
     @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.num_allies, ally_key)
-        team_1_start = spawn_fn(*self.low_sector, self.num_enemies, enemy_key)
+        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)
         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)
         unit_weapon_cooldowns = jnp.zeros((self.num_agents,))
         # default behaviour spawn all marines
-        unit_types = self.scenario.unit_types
+        unit_types = cast(Array, self.scenario.unit_types)
         unit_health = self.unit_type_health[unit_types]
         state = State(
             unit_positions=unit_positions,
@@ -109,7 +122,7 @@ class Environment(SMAX):
             terminal=False,
             unit_weapon_cooldowns=unit_weapon_cooldowns,
         )
-        state = self._push_units_away(state)
+        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)
@@ -141,7 +154,7 @@ class Environment(SMAX):
         key: chex.PRNGKey,
         state: State,
         actions: Tuple[chex.Array, chex.Array],
-    ) -> Tuple[Dict[str, chex.Array], State, Dict[str, float], Dict[str, bool], Dict]:
+    ) -> State:
         @partial(jax.vmap, in_axes=(None, None, 0, 0))
         def intersect_fn(pos, new_pos, obs, obs_end):
             d1 = jnp.cross(obs - pos, new_pos - pos)
@@ -167,7 +180,7 @@ class Environment(SMAX):
             # 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]
+            pos = cast(Array, state.unit_positions[idx])
             new_pos = (
                 pos
                 + vec
@@ -219,6 +232,7 @@ class Environment(SMAX):
                 lambda: action + self.num_allies - self.num_movement_actions,
                 lambda: self.num_allies - 1 - (action - self.num_movement_actions),
             )
+            attacked_idx = cast(int, attacked_idx)  # Cast to int
             # 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
@@ -250,7 +264,7 @@ class Environment(SMAX):
             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.bool_) # type: ignore
                 .astype(jnp.float32)
             )
             bystander_health_diff = (
@@ -365,7 +379,8 @@ class Environment(SMAX):
         #########################################################
 
         unit_weapon_cooldowns = state.unit_weapon_cooldowns + cooldown_diff
-        state = state.replace(
+        # replace unit health, unit positions and unit weapon cooldowns
+        state = state.replace(  # type: ignore
             unit_health=unit_health,
             unit_positions=pos,
             unit_weapon_cooldowns=unit_weapon_cooldowns,

parabellum/gun.py

@@ -0,0 +1,70 @@
+# gun.py
+#   parabellum bullet rendering assosciated functions
+# by: Noah Syrkis
+
+# imports
+from functools import partial
+import jax.numpy as jnp
+
+
+def dist_fn(env, pos):  # computing the distances between all ally and enemy agents
+    delta = pos[None, :, :] - pos[:, None, :]
+    dist = jnp.sqrt((delta**2).sum(axis=2))
+    dist = dist[: env.num_allies, env.num_allies :]
+    return {"ally": dist, "enemy": dist.T}
+
+
+def range_fn(env, dists, ranges):  # computing what targets are in range
+    ally_range = dists["ally"] < ranges[: env.num_allies][:, None]
+    enemy_range = dists["enemy"] < ranges[env.num_allies :][:, None]
+    return {"ally": ally_range, "enemy": enemy_range}
+
+
+def target_fn(acts, in_range, team):  # computing the one hot valid targets
+    t_acts = jnp.stack([v for k, v in acts.items() if k.startswith(team)]).T
+    t_targets = jnp.where(t_acts > 4, -1, t_acts - 5)  # first 5 are move actions
+    t_attacks = jnp.eye(in_range[team].shape[1] + 1)[t_targets][:, :-1]
+    return t_attacks * in_range[team]  # one hot valid targets
+
+
+def attack_fn(env, state_seq):  # one hot attack list
+    attacks = []
+    for _, state, acts in state_seq:
+        dists = dist_fn(env, state.unit_positions)
+        ranges = env.unit_type_attack_ranges[state.unit_types]
+        in_range = range_fn(env, dists, ranges)
+        target = partial(target_fn, acts, in_range)
+        attack = {"ally": target("ally"), "enemy": target("enemy")}
+        attacks.append(attack)
+    return attacks
+
+
+def bullet_fn(env, states):
+    bullet_seq = []
+    attack_seq = attack_fn(env, states)
+
+    def aux_fn(team):
+        bullets = jnp.stack(jnp.where(one_hot[team] == 1)).T
+        # bullets = bullets.at[:, 2 if team == "ally" else 1].add(env.num_allies)
+        return bullets
+
+    state_zip = zip(states[:-1], states[1:])
+    for i, ((_, state, _), (_, n_state, _)) in enumerate(state_zip):
+        one_hot = attack_seq[i]
+        ally_bullets, enemy_bullets = aux_fn("ally"), aux_fn("enemy")
+
+        ally_bullets_source = state.unit_positions[ally_bullets[:, 0]]
+        enemy_bullets_target = n_state.unit_positions[enemy_bullets[:, 1]]
+
+        enemy_bullets_source = state.unit_positions[
+            enemy_bullets[:, 0] + env.num_allies
+        ]
+        ally_bullets_target = n_state.unit_positions[
+            ally_bullets[:, 1] + env.num_allies
+        ]
+
+        ally_bullets = jnp.stack((ally_bullets_source, ally_bullets_target), axis=1)
+        enemy_bullets = jnp.stack((enemy_bullets_source, enemy_bullets_target), axis=1)
+
+        bullet_seq.append((ally_bullets, enemy_bullets))
+    return bullet_seq

parabellum/map.py

@@ -1,5 +1,5 @@
 # map.py
-#   parabellum map functions
+    # parabellum map functions
 # by: Noah Syrkis
 
 # imports
@@ -7,43 +7,62 @@ import jax.numpy as jnp
 from geopy.geocoders import Nominatim
 import geopandas as gpd
 import osmnx as ox
-import rasterio
-from jax import random
+import contextily as cx
+import matplotlib.pyplot as plt
 from rasterio import features
 import rasterio.transform
+from typing import Optional, Tuple
+from geopy.location import Location
+from shapely.geometry import Point
 
 # constants
 geolocator = Nominatim(user_agent="parabellum")
-tags = {"building": True}
+BUILDING_TAGS = {"building": True}
 
+def get_location(place: str) -> Tuple[float, float]:
+    """Get coordinates for a given place."""
+    coords: Optional[Location] = geolocator.geocode(place)  # type: ignore
+    if coords is None:
+        raise ValueError(f"Could not geocode the place: {place}")
+    return (coords.latitude, coords.longitude)
 
-# functions
-def terrain_fn(place: str, size: int = 1000):
-    """Returns a rasterized map of a given location."""
+def get_building_geometry(point: Tuple[float, float], size: int) -> gpd.GeoDataFrame:
+    """Get building geometry for a given point and size."""
+    geometry = ox.features_from_point(point, tags=BUILDING_TAGS, dist=size // 2)
+    return gpd.GeoDataFrame(geometry).set_crs("EPSG:4326")
 
-    # location info
-    location = geolocator.geocode(place)
-    coords = (location.latitude, location.longitude)
+def rasterize_geometry(gdf: gpd.GeoDataFrame, size: int) -> jnp.ndarray:
+    """Rasterize geometry and return as a JAX array."""
+    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)
 
-    # shape info
-    geometry = ox.features_from_point(coords, tags=tags, dist=size // 2)
-    gdf = gpd.GeoDataFrame(geometry).set_crs("EPSG:4326")
+def terrain_fn(place: str, size: int = 1000) -> 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)
+    return mask, base
 
-    # raster info
-    t = rasterio.transform.from_bounds(*gdf.total_bounds, size, size)
-    raster = features.rasterize(gdf.geometry, out_shape=(size, size), transform=t)
-
-    # rotate 180 degrees
-    raster = jnp.rot90(raster, 2)
-
-    return jnp.array(raster).astype(jnp.uint8)
+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)
+    gdf = get_building_geometry(point, size)
+    basemap, _ = cx.bounds2img(*gdf.total_bounds, ll=True)
+    # 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)
 
 
 if __name__ == "__main__":
     import seaborn as sns
+    place = "Thun, Switzerland"
+    mask, base = terrain_fn(place)
 
-    place = "Vesterbro, Copenhagen, Denmark"
-    terrain = terrain_fn(place)
-    rng, key = random.split(random.PRNGKey(0))
-    agents = spawn_fn(terrain, 12, 100, key)
-    print(agents)
+    fig, ax = plt.subplots(1, 2, figsize=(10, 5))
+    ax[0].imshow(mask) # type: ignore
+    ax[1].imshow(base) # type: ignore
+    plt.show()

parabellum/run.py

@@ -10,6 +10,7 @@ import darkdetect
 import jax.numpy as jnp
 from chex import dataclass
 import jaxmarl
+from jax import Array
 from typing import Tuple, List, Dict, Optional
 import parabellum as pb
 
@@ -20,7 +21,7 @@ bg = (0, 0, 0) if darkdetect.isDark() else (255, 255, 255)
 
 
 # types
-State = jaxmarl.environments.smax.smax_env.State
+State = jaxmarl.environments.smax.smax_env.State # type: ignore
 Obs = Reward = Done = Action = Dict[str, jnp.ndarray]
 StateSeq = List[Tuple[jnp.ndarray, State, Action]]
 
@@ -35,12 +36,12 @@ class Control:
 @dataclass
 class Game:
     clock: pygame.time.Clock
-    state: State
+    state: State  # type: ignore
     obs: Dict
     state_seq: StateSeq
     control: Control
     env: pb.Environment
-    rng: random.PRNGKey
+    rng: Array
 
 
 def handle_event(event, control_state):
@@ -100,7 +101,7 @@ def step_fn(game):
 
 # state
 if __name__ == "__main__":
-    env = pb.Parabellum(pb.scenarios["default"])
+    env = pb.Environment(pb.scenarios["default"])
     pygame.init()
     screen = pygame.display.set_mode((1000, 1000))
     render = partial(render_fn, screen)
@@ -115,7 +116,7 @@ if __name__ == "__main__":
         state=state,
         obs=obs,
     )
-    game = Game(**kwargs)
+    game = Game(**kwargs)  # type: ignore
 
     while game.control.running:
         game = control_fn(game)

parabellum/vis.py

@@ -2,36 +2,131 @@
 Visualizer for the Parabellum environment
 """
 
-from tqdm import tqdm
-import jax.numpy as jnp
-import jax
-from jax import vmap
-from jax import tree_util
+# Standard library imports
 from functools import partial
-import darkdetect
-import numpy as np
-import pygame
-import matplotlib.pyplot as plt
-import os
-from moviepy.editor import ImageSequenceClip
-from typing import Optional
+from typing import Optional, List, Tuple
+import cv2
+from PIL import Image
+
+# JAX and JAX-related imports
+import jax
+from chex import dataclass
+import chex
+from jax import vmap, tree_util, Array, jit
+import jax.numpy as jnp
 from jaxmarl.environments.multi_agent_env import MultiAgentEnv
+from jaxmarl.environments.smax import SMAX
 from jaxmarl.viz.visualizer import SMAXVisualizer
 
-# default dict
-from collections import defaultdict
+# Third-party imports
+import numpy as np
+import pygame
+import cv2
+from tqdm import tqdm
+
+# Local imports
+import parabellum as pb
+
+
+# skin dataclass
+@dataclass
+class Skin:
+    # basemap: Array  # basemap of buildings
+    maskmap: Array  # maskmap of buildings
+    bg: Tuple[int, int, int] = (255, 255, 255)
+    fg: Tuple[int, int, int] = (0, 0, 0)
+    ally: Tuple[int, int, int]  = (0, 255, 0)
+    enemy: Tuple[int, int, int]  = (255, 0, 0)
+    pad: int  = 100
+    size: int  = 1000  # excluding padding
+    fps: int = 24
+    vis_size: int = 1000   # size of the map in Vis (exluding padding)
+    scale: Optional[float] = None
 
 
-# constants
-action_to_symbol = {0: "↑", 1: "→", 2: "↓", 3: "←", 4: "Ø"}
+class Visualizer(SMAXVisualizer):
+    def __init__(self, env: pb.Environment, state_seq, skin: Skin, reward_seq=None):
+        super(Visualizer, self).__init__(env, state_seq, reward_seq)
 
+        # self.bullet_seq = vmap(partial(bullet_fn, self.env))(self.state_seq)
+        self.action_seq = [action for _, _, action in state_seq]  # bcs SMAX bug
+        self.state_seq = state_seq
+        self.image = image_fn(skin)
+        self.skin = skin
+        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)
+        for idx, (state_seq, action_seq) in enumerate(zip(state_seq_seq, action_seq_seq)):
+            animate_fn(self.env, self.skin, self.image, state_seq, action_seq, f"{save_fname}_{idx}.mp4")
+
+
+# functions
+def animate_fn(env, skin, image, state_seq, action_seq, save_fname):
+    pygame.init()
+    frames = []
+    for idx, (state_tup, action) in enumerate(zip(state_seq, action_seq)):
+        frames += [frame_fn(env, skin, image, state_tup[1], action, idx)]
+    # use cv2 to write frames to video
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # type: ignore
+    out = cv2.VideoWriter(save_fname, fourcc, skin.fps, (skin.size + skin.pad * 2, skin.size + skin.pad * 2))
+    for frame in frames:
+        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    out.release()
+    pygame.quit()
 
-def small_multiples():
-    # make video of small multiples based on all videos in output
-    video_files = [f"output/parabellum_{i}.mp4" for i in range(4)]
-    # load mp4 videos and make a grid
-    clips = [ImageSequenceClip.load(filename) for filename in video_files]
-    print(len(clips))
+
+def init_frame(env, skin, image, state: pb.State, action: Array, idx: int) -> pygame.Surface:
+    dims = (skin.size + skin.pad * 2, skin.size + skin.pad * 2)
+    frame = pygame.Surface(dims, pygame.SRCALPHA | pygame.HWSURFACE)
+    return frame
+
+
+def transform_frame(env, skin, frame):
+    frame = np.rot90(pygame.surfarray.pixels3d(frame).swapaxes(0, 1), 2)
+    return frame
+
+
+def frame_fn(env, skin, image, state: pb.State, action: Array, idx: int) -> np.ndarray:
+    """Create a frame"""
+    frame = init_frame(env, skin, image, state, action, idx)
+
+    pipeline = [render_background, render_agents, render_action, render_bullet]
+    for fn in pipeline:
+        frame = fn(env, skin, image, frame, state, action)
+
+    return transform_frame(env, skin, frame)
+
+
+def render_background(env, skin, image, frame, state, action):
+    coords = (skin.pad-5, skin.pad-5, skin.size+10, skin.size+10)
+    frame.fill(skin.bg)
+    frame.blit(image, coords)
+    pygame.draw.rect(frame, skin.fg, coords, 3)
+    return frame
+
+
+def render_action(env, skin, image, frame, state, action):
+    return frame
+
+
+def render_bullet(env, skin, image, frame, state, action):
+    return frame
+
+def render_agents(env, skin, image, frame, state, action):
+    units = state.unit_positions, state.unit_teams, state.unit_types, state.unit_health
+    for idx, (pos, team, kind, health) in enumerate(zip(*units)):
+        pos = tuple((pos * skin.scale).astype(int) + skin.pad)
+        # draw the agent
+        if health > 0:
+            unit_size = env.unit_type_radiuses[kind]
+            radius = float(jnp.ceil((unit_size * skin.scale)).astype(int) + 1)
+            pygame.draw.circle(frame, skin.fg, pos, radius, 1)
+            pygame.draw.circle(frame, skin.bg, pos, radius + 1, 1)
+    return frame
 
 
 def text_fn(text):
@@ -39,263 +134,36 @@ def text_fn(text):
     return pygame.transform.rotate(text, 180)
 
 
-class Visualizer(SMAXVisualizer):
-    def __init__(self, env: MultiAgentEnv, state_seq, reward_seq=None, skin=None):
-        self.fig, self.ax = None, None
-        super().__init__(env, state_seq, reward_seq)
-        # clear the figure made by SMAXVisualizer
-        plt.close()
-        self.bg = (0, 0, 0) if darkdetect.isDark() else (255, 255, 255)
-        self.fg = (255, 255, 255) if darkdetect.isDark() else (0, 0, 0)
-        self.pad = 75
-        # TODO: make sure it's always a 1024x1024 image
-        self.width = 1000
-        self.s = self.width + self.pad + self.pad
-        self.scale = self.width / env.map_width
-        self.action_seq = [action for _, _, action in state_seq]  # bcs SMAX bug
-        # self.bullet_seq = vmap(partial(bullet_fn, self.env))(self.state_seq)
+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 = (motif > 0).astype(np.uint8)
+    image = np.zeros((skin.size, skin.size, 3), dtype=np.uint8) + skin.bg
+    image[motif == 1] = skin.fg
+    image = pygame.surfarray.make_surface(image)
+    image = pygame.transform.scale(image, (skin.size, skin.size))
+    return image
 
-    def animate(self, save_fname: str = "output/parabellum.mp4"):
-        multi_dim = self.state_seq[0][1].unit_positions.ndim > 2
-        if multi_dim:
-            n_envs = self.state_seq[0][1].unit_positions.shape[0]
-            if not self.have_expanded:
-                state_seqs = vmap(self.env.expand_state_seq)(self.state_seq)
-                self.have_expanded = True
-            for i in range(n_envs):
-                state_seq = jax.tree_map(lambda x: x[i], state_seqs)
-                action_seq = jax.tree_map(lambda x: x[i], self.action_seq)
-                self.animate_one(
-                    state_seq, action_seq, save_fname.replace(".mp4", f"_{i}.mp4")
-                )
-        else:
-            state_seq = self.env.expand_state_seq(self.state_seq)
-            self.animate_one(state_seq, self.action_seq, save_fname)
-
-    def animate_one(self, state_seq, action_seq, save_fname):
-        frames = []  # frames for the video
-        pygame.init()  # initialize pygame
-        terrain = np.array(self.env.terrain_raster.T)
-        rgb_array = np.zeros((terrain.shape[0], terrain.shape[1], 3), dtype=np.uint8)
-        if darkdetect.isLight():
-            rgb_array += 255
-        rgb_array[terrain == 1] = self.fg
-        mask_surface = pygame.surfarray.make_surface(rgb_array)
-        mask_surface = pygame.transform.scale(mask_surface, (self.width, self.width))
-
-        for idx, (_, state, _) in tqdm(enumerate(state_seq), total=len(self.state_seq)):
-            action = action_seq[idx // self.env.world_steps_per_env_step]
-            screen = pygame.Surface(
-                (self.s, self.s), pygame.HWSURFACE | pygame.DOUBLEBUF
-            )
-            screen.fill(self.bg)  # fill the screen with the background color
-            screen.blit(
-                mask_surface,
-                (self.pad, self.pad, self.width, self.width),
-            )
-            # add env.scenario.place to the title (in top padding)
-            font = pygame.font.SysFont("Fira Code", 18)
-            width = self.env.map_width
-            title = f"{width}x{width}m in {self.env.scenario.place}"
-            text = text_fn(font.render(title, True, self.fg))
-            # center the text
-            screen.blit(text, (self.s // 2 - text.get_width() // 2, self.pad // 4))
-            # draw edge around terrain
-            pygame.draw.rect(
-                screen,
-                self.fg,
-                (
-                    self.pad - 2,
-                    self.pad - 2,
-                    self.width + 4,
-                    self.width + 4,
-                ),
-                2,
-            )
-
-            self.render_agents(screen, state)  # render the agents
-            self.render_action(screen, action)
-            # self.render_obstacles(screen)  # render the obstacles
-
-            # bullets
-            """ if idx < len(self.bullet_seq) * 8:
-                bullets = self.bullet_seq[idx // 8]
-                self.render_bullets(screen, bullets, idx % 8) """
-
-            # rotate the screen and append to frames
-            pixels = pygame.surfarray.pixels3d(screen).swapaxes(0, 1)
-            # rotate the screen 180 degrees (transpose and flip)
-            pixels = np.rot90(pixels, 2)  # pygame starts in bottom left
-            frames.append(pixels)
-        # save the images
-        clip = ImageSequenceClip(frames, fps=48)
-        clip.write_videofile(save_fname, fps=48)
-        clip.write_gif(save_fname.replace(".mp4", ".gif"), fps=24)
-        pygame.quit()
-
-    def render_agents(self, screen, state):
-        time_tuple = zip(
-            state.unit_positions,
-            state.unit_teams,
-            state.unit_types,
-            state.unit_health,
-        )
-        for idx, (pos, team, kind, hp) in enumerate(time_tuple):
-            face_col = self.fg if int(team.item()) == 0 else self.bg
-            pos = tuple(((pos * self.scale) + self.pad).tolist())
-            # draw the agent
-            if hp > 0:
-                hp_frac = hp / self.env.unit_type_health[kind]
-                unit_size = self.env.unit_type_radiuses[kind]
-                radius = jnp.ceil((unit_size * self.scale * hp_frac)).astype(int) + 1
-                pygame.draw.circle(screen, face_col, pos, radius)
-                pygame.draw.circle(screen, self.fg, pos, radius, 1)
-
-            # draw the sight range
-            # sight_range = self.env.unit_type_sight_ranges[kind] * self.scale
-            # pygame.draw.circle(screen, self.fg, pos, sight_range.astype(int), 2)
-
-            # draw attack range
-            # attack_range = self.env.unit_type_attack_ranges[kind] * self.scale
-            # pygame.draw.circle(screen, self.fg, pos, attack_range.astype(int), 2)
-            # work out which agents are being shot
-
-    def render_action(self, screen, action):
-        if self.env.action_type != "discrete":
-            return
-
-        def coord_fn(idx, n, team, text):
-            text_adj = text.get_width() / 2
-            is_ally = team == "ally"
-            return (
-                # vertically centered so that n / 2 is above and below the center
-                self.pad + self.width + self.pad / 2 - text_adj
-                if is_ally
-                else self.pad / 2 - text_adj,
-                self.s / 2 - (n / 2) * self.s / 20 + idx * self.s / 20,
-            )
-
-        for team, number in [("ally", 0), ("enemy", 1)]:
-            for idx in range(self.env.num_allies):
-                symb = action_to_symbol.get(
-                    action[f"{team}_{idx}"].astype(int).item(), "Ø"
-                )
-                font = pygame.font.SysFont(
-                    "Fira Code", jnp.sqrt(self.s).astype(int).item()
-                )
-                text = text_fn(font.render(symb, True, self.fg))
-                coord = coord_fn(idx, self.env.num_allies, team, text)
-                screen.blit(text, coord)
-
-    def render_obstacles(self, screen):
-        for c, d in zip(self.env.obstacle_coords, self.env.obstacle_deltas):
-            d = tuple(((c + d) * self.scale).tolist())
-            c = tuple((c * self.scale).tolist())
-            pygame.draw.line(screen, self.fg, c, d, 5)
-
-    def render_bullets(self, screen, bullets, jdx):
-        jdx += 1
-        ally_bullets, enemy_bullets = bullets
-        for source, target in ally_bullets:
-            position = source + (target - source) * jdx / 8
-            position *= self.scale
-            pygame.draw.circle(screen, self.fg, tuple(position.tolist()), 3)
-        for source, target in enemy_bullets:
-            position = source + (target - source) * jdx / 8
-            position *= self.scale
-            pygame.draw.circle(screen, self.fg, tuple(position.tolist()), 3)
 
+def unbatch_fn(state_seq, action_seq):
+    """state seq is a list of tuples of (step_key, state, actions)."""
+    if is_multi_run(state_seq):
+        n_envs = state_seq[0][1].unit_positions.shape[0]
+        state_seq_seq = [jax.tree_map(lambda x: x[i], state_seq) for i in range(n_envs)]
+        action_seq_seq = [jax.tree_map(lambda x: x[i], action_seq) for i in range(n_envs)]
+    else:
+        state_seq_seq = [state_seq]
+        action_seq_seq = [action_seq]
+    return state_seq_seq, action_seq_seq
 
-# functions
-# bullet functions
-def dist_fn(env, pos):  # computing the distances between all ally and enemy agents
-    delta = pos[None, :, :] - pos[:, None, :]
-    dist = jnp.sqrt((delta**2).sum(axis=2))
-    dist = dist[: env.num_allies, env.num_allies :]
-    return {"ally": dist, "enemy": dist.T}
-
-
-def range_fn(env, dists, ranges):  # computing what targets are in range
-    ally_range = dists["ally"] < ranges[: env.num_allies][:, None]
-    enemy_range = dists["enemy"] < ranges[env.num_allies :][:, None]
-    return {"ally": ally_range, "enemy": enemy_range}
-
-
-def target_fn(acts, in_range, team):  # computing the one hot valid targets
-    t_acts = jnp.stack([v for k, v in acts.items() if k.startswith(team)]).T
-    t_targets = jnp.where(t_acts > 4, -1, t_acts - 5)  # first 5 are move actions
-    t_attacks = jnp.eye(in_range[team].shape[1] + 1)[t_targets][:, :-1]
-    return t_attacks * in_range[team]  # one hot valid targets
-
-
-def attack_fn(env, state_seq):  # one hot attack list
-    attacks = []
-    for _, state, acts in state_seq:
-        dists = dist_fn(env, state.unit_positions)
-        ranges = env.unit_type_attack_ranges[state.unit_types]
-        in_range = range_fn(env, dists, ranges)
-        target = partial(target_fn, acts, in_range)
-        attack = {"ally": target("ally"), "enemy": target("enemy")}
-        attacks.append(attack)
-    return attacks
-
-
-def bullet_fn(env, states):
-    bullet_seq = []
-    attack_seq = attack_fn(env, states)
-
-    def aux_fn(team):
-        bullets = jnp.stack(jnp.where(one_hot[team] == 1)).T
-        # bullets = bullets.at[:, 2 if team == "ally" else 1].add(env.num_allies)
-        return bullets
-
-    state_zip = zip(states[:-1], states[1:])
-    for i, ((_, state, _), (_, n_state, _)) in enumerate(state_zip):
-        one_hot = attack_seq[i]
-        ally_bullets, enemy_bullets = aux_fn("ally"), aux_fn("enemy")
-
-        ally_bullets_source = state.unit_positions[ally_bullets[:, 0]]
-        enemy_bullets_target = n_state.unit_positions[enemy_bullets[:, 1]]
-
-        enemy_bullets_source = state.unit_positions[
-            enemy_bullets[:, 0] + env.num_allies
-        ]
-        ally_bullets_target = n_state.unit_positions[
-            ally_bullets[:, 1] + env.num_allies
-        ]
-
-        ally_bullets = jnp.stack((ally_bullets_source, ally_bullets_target), axis=1)
-        enemy_bullets = jnp.stack((enemy_bullets_source, enemy_bullets_target), axis=1)
-
-        bullet_seq.append((ally_bullets, enemy_bullets))
-    return bullet_seq
-
-
-# test the visualizer
-if __name__ == "__main__":
-    from jax import random, numpy as jnp
-    from parabellum import Parabellum, scenarios
-
-    # small_multiples()  # testing small multiples (not working yet)
-    # exit()
-
-    n_envs = 2
-    env = Parabellum(scenarios["default"], action_type="discrete")
-    rng, reset_rng = random.split(random.PRNGKey(0))
-    reset_key = random.split(reset_rng, n_envs)
-    obs, state = vmap(env.reset)(reset_key)
-    state_seq = []
-
-    for i in range(100):
-        rng, act_rng, step_rng = random.split(rng, 3)
-        act_key = random.split(act_rng, (len(env.agents), n_envs))
-        act = {
-            a: vmap(env.action_space(a).sample)(act_key[i])
-            for i, a in enumerate(env.agents)
-        }
-        step_key = random.split(step_rng, n_envs)
-        state_seq.append((step_key, state, act))
-        obs, state, reward, done, infos = vmap(env.step)(step_key, state, act)
-
-    vis = Visualizer(env, state_seq)
-    vis.animate()
+
+def expand_fn(env, state_seq, action_seq):
+    """Expand the state sequence"""
+    fn = env.expand_state_seq
+    state_seq = vmap(fn)(state_seq) if is_multi_run(state_seq) else fn(state_seq)
+    action_seq = [action_seq[i // env.world_steps_per_env_step] for i in range(len(state_seq))]
+    return state_seq, action_seq
+
+
+def is_multi_run(state_seq):
+    return state_seq[0][1].unit_positions.ndim > 2

{parabellum-0.2.20.dist-info → parabellum-0.2.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: parabellum
-Version: 0.2.20
+Version: 0.2.21
 Summary: Parabellum environment for parallel warfare simulation
 Home-page: https://github.com/syrkis/parabellum
 License: MIT
@@ -12,15 +12,20 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: contextily (>=1.6.0,<2.0.0)
 Requires-Dist: darkdetect (>=0.8.0,<0.9.0)
+Requires-Dist: folium (>=0.17.0,<0.18.0)
 Requires-Dist: geopandas (>=1.0.0,<2.0.0)
 Requires-Dist: geopy (>=2.4.1,<3.0.0)
+Requires-Dist: ipykernel (>=6.29.5,<7.0.0)
 Requires-Dist: jax (==0.4.17)
 Requires-Dist: jaxmarl (==0.0.3)
 Requires-Dist: jupyterlab (>=4.2.2,<5.0.0)
 Requires-Dist: moviepy (>=1.0.3,<2.0.0)
 Requires-Dist: numpy (<2)
+Requires-Dist: opencv-python (>=4.10.0.84,<5.0.0.0)
 Requires-Dist: osmnx (>=1.9.3,<2.0.0)
+Requires-Dist: pandas (>=2.2.2,<3.0.0)
 Requires-Dist: poetry (>=1.8.3,<2.0.0)
 Requires-Dist: pygame (>=2.5.2,<3.0.0)
 Requires-Dist: rasterio (>=1.3.10,<2.0.0)

parabellum-0.2.21.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=L6GHlLxywpkV1bRnZcYBURREPP4CRfet_pEwCt5DB04,16724
+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.21.dist-info/METADATA,sha256=-K-3eYl1BvR3tFsiTxTyfHErQJdgPQZx08iq_kS2544,2671
+parabellum-0.2.21.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.2.21.dist-info/RECORD,,

parabellum-0.2.20.dist-info/RECORD

@@ -1,8 +0,0 @@
-parabellum/__init__.py,sha256=cI1kxVQ274VZrBLxUYvB_j6UqhFY44sId4NgVjbf678,245
-parabellum/env.py,sha256=B2WaEzCQnakVB1AuYqFsb7aiPIbaUsuq0CjdKq4pQm8,16204
-parabellum/map.py,sha256=CdPebRuGafj9fAydAyTa8bqpr-PldZw7heKRReGTDBg,1257
-parabellum/run.py,sha256=HyPdz5iVD8q0iYaZL2Nf02fGHByHpecUGoPlQrq9v8s,3411
-parabellum/vis.py,sha256=uIhN9VSJlT4XoNuNCiU_Bw2VSNoJFOYlScCTNrMkGsg,11977
-parabellum-0.2.20.dist-info/METADATA,sha256=IO3Wsx_nTM3v63FtbnJqaEfPfXOyhNGtKOm_ppchUqo,2453
-parabellum-0.2.20.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-parabellum-0.2.20.dist-info/RECORD,,