parabellum 0.0.0__py3-none-any.whl

parabellum/.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1,4 @@
+from .env import Parabellum, Scenario
+from .vis import Visualizer
+
+__all__ = ["Parabellum", "Visualizer", "Scenario"]

parabellum/.ipynb_checkpoints/env-checkpoint.py

@@ -0,0 +1,296 @@
+"""Parabellum environment based on SMAX"""
+
+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
+from functools import partial
+
+
+@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
+
+            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]]
+            )
+
+            #########################################################
+            #########################################################
+
+            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 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
+
+        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)
+
+        # 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
+
+        _, bystander_health_diff = bystander
+        unit_health -= bystander_health_diff.sum(axis=0)  # might be axis=1
+
+        #########################################################
+        #########################################################
+
+        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))
+
+

parabellum/.ipynb_checkpoints/vis-checkpoint.py

@@ -0,0 +1,230 @@
+"""Visualizer for the Parabellum environment"""
+
+from tqdm import tqdm
+import jax.numpy as jnp
+import jax
+from jax import vmap
+from functools import partial
+import darkdetect
+import pygame
+from moviepy.editor import ImageSequenceClip
+from typing import Optional
+from jaxmarl.environments.multi_agent_env import MultiAgentEnv
+from jaxmarl.viz.visualizer import SMAXVisualizer
+
+# default dict
+from collections import defaultdict
+
+
+# constants
+action_to_symbol = {0: "↑", 1: "→", 2: "↓", 3: "←", 4: "Ø"}
+
+
+class Visualizer(SMAXVisualizer):
+    def __init__(self, env: MultiAgentEnv, state_seq, reward_seq=None):
+        super().__init__(env, state_seq, reward_seq)
+        # remove fig and ax from super
+        self.fig, self.ax = None, None
+        self.bg = (0, 0, 0) if darkdetect.isDark() else (255, 255, 255)
+        self.fg = (235, 235, 235) if darkdetect.isDark() else (20, 20, 20)
+        self.s = 1000
+        self.scale = self.s / self.env.map_width
+        self.action_seq = [action for _, _, action in state_seq]  # bcs SMAX bug
+        self.bullet_seq = bullet_fn(self.env, self.state_seq)
+
+    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).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):
+        def coord_fn(idx, n, team):
+            return (
+                self.s / 20 if team == 0 else self.s - self.s / 20,
+                # vertically centered so that n / 2 is above and below the center
+                self.s / 2 - (n / 2) * self.s / 20 + idx * self.s / 20,
+            )
+
+        for idx in range(self.env.num_allies):
+            symb = action_to_symbol.get(action[f"ally_{idx}"].astype(int).item(), "Ø")
+            font = pygame.font.SysFont("Fira Code", jnp.sqrt(self.s).astype(int).item())
+            text = font.render(symb, True, self.fg)
+            coord = coord_fn(idx, self.env.num_allies, 0)
+            screen.blit(text, coord)
+
+        for idx in range(self.env.num_enemies):
+            symb = action_to_symbol.get(action[f"enemy_{idx}"].astype(int).item(), "Ø")
+            font = pygame.font.SysFont("Fira Code", jnp.sqrt(self.s).astype(int).item())
+            text = font.render(symb, True, self.fg)
+            coord = coord_fn(idx, self.env.num_enemies, 1)
+            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 animate(self, save_fname: str = "parabellum.mp4"):
+        if not self.have_expanded:
+            self.expand_state_seq()
+        frames = []  # frames for the video
+        pygame.init()  # initialize pygame
+        for idx, (_, state, _) in tqdm(
+            enumerate(self.state_seq), total=len(self.state_seq)
+        ):
+            screen = pygame.Surface(
+                (self.s, self.s), pygame.HWSURFACE | pygame.DOUBLEBUF
+            )
+            screen.fill(self.bg)  # fill the screen with the background color
+
+            self.render_agents(screen, state)  # render the agents
+            self.render_action(screen, self.action_seq[idx // 8])
+            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
+            frames.append(pygame.surfarray.pixels3d(screen).swapaxes(0, 1))
+
+        # 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()
+
+        return clip
+
+
+# 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 parabellum import Parabellum, Scenario
+    from jax import random, numpy as jnp
+
+    s = Scenario(jnp.array([[16, 0]]),
+                jnp.array([[0, 32]]) * 8,
+                jnp.zeros((19,), dtype=jnp.uint8),
+                9,
+                 10)
+    env = Parabellum(map_width=32, map_height=32, walls_cause_death=False, scenario=s)
+    rng, key = random.split(random.PRNGKey(0))
+    obs, state = env.reset(key)
+    state_seq = []
+    for step in range(50):
+        rng, key = random.split(rng)
+        key_act = random.split(key, len(env.agents))
+        actions = {
+            agent: jnp.array(1)
+            for i, agent in enumerate(env.agents)
+        }
+        state_seq.append((key, state, actions))
+        rng, key_step = random.split(rng)
+        obs, state, reward, done, infos = env.step(key_step, state, actions)
+
+    vis = Visualizer(env, state_seq)
+    vis.animate()
+
+

parabellum/__init__.py

@@ -0,0 +1,4 @@
+from .env import Parabellum, Scenario
+from .vis import Visualizer
+
+__all__ = ["Parabellum", "Visualizer", "Scenario"]

parabellum/env.py

@@ -0,0 +1,296 @@
+"""Parabellum environment based on SMAX"""
+
+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
+from functools import partial
+
+
+@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
+
+            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]]
+            )
+
+            #########################################################
+            #########################################################
+
+            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 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
+
+        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)
+
+        # 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
+
+        _, bystander_health_diff = bystander
+        unit_health -= bystander_health_diff.sum(axis=0)  # might be axis=1
+
+        #########################################################
+        #########################################################
+
+        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))
+
+

parabellum/vis.py

@@ -0,0 +1,230 @@
+"""Visualizer for the Parabellum environment"""
+
+from tqdm import tqdm
+import jax.numpy as jnp
+import jax
+from jax import vmap
+from functools import partial
+import darkdetect
+import pygame
+from moviepy.editor import ImageSequenceClip
+from typing import Optional
+from jaxmarl.environments.multi_agent_env import MultiAgentEnv
+from jaxmarl.viz.visualizer import SMAXVisualizer
+
+# default dict
+from collections import defaultdict
+
+
+# constants
+action_to_symbol = {0: "↑", 1: "→", 2: "↓", 3: "←", 4: "Ø"}
+
+
+class Visualizer(SMAXVisualizer):
+    def __init__(self, env: MultiAgentEnv, state_seq, reward_seq=None):
+        super().__init__(env, state_seq, reward_seq)
+        # remove fig and ax from super
+        self.fig, self.ax = None, None
+        self.bg = (0, 0, 0) if darkdetect.isDark() else (255, 255, 255)
+        self.fg = (235, 235, 235) if darkdetect.isDark() else (20, 20, 20)
+        self.s = 1000
+        self.scale = self.s / self.env.map_width
+        self.action_seq = [action for _, _, action in state_seq]  # bcs SMAX bug
+        self.bullet_seq = bullet_fn(self.env, self.state_seq)
+
+    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).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):
+        def coord_fn(idx, n, team):
+            return (
+                self.s / 20 if team == 0 else self.s - self.s / 20,
+                # vertically centered so that n / 2 is above and below the center
+                self.s / 2 - (n / 2) * self.s / 20 + idx * self.s / 20,
+            )
+
+        for idx in range(self.env.num_allies):
+            symb = action_to_symbol.get(action[f"ally_{idx}"].astype(int).item(), "Ø")
+            font = pygame.font.SysFont("Fira Code", jnp.sqrt(self.s).astype(int).item())
+            text = font.render(symb, True, self.fg)
+            coord = coord_fn(idx, self.env.num_allies, 0)
+            screen.blit(text, coord)
+
+        for idx in range(self.env.num_enemies):
+            symb = action_to_symbol.get(action[f"enemy_{idx}"].astype(int).item(), "Ø")
+            font = pygame.font.SysFont("Fira Code", jnp.sqrt(self.s).astype(int).item())
+            text = font.render(symb, True, self.fg)
+            coord = coord_fn(idx, self.env.num_enemies, 1)
+            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 animate(self, save_fname: str = "parabellum.mp4"):
+        if not self.have_expanded:
+            self.expand_state_seq()
+        frames = []  # frames for the video
+        pygame.init()  # initialize pygame
+        for idx, (_, state, _) in tqdm(
+            enumerate(self.state_seq), total=len(self.state_seq)
+        ):
+            screen = pygame.Surface(
+                (self.s, self.s), pygame.HWSURFACE | pygame.DOUBLEBUF
+            )
+            screen.fill(self.bg)  # fill the screen with the background color
+
+            self.render_agents(screen, state)  # render the agents
+            self.render_action(screen, self.action_seq[idx // 8])
+            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
+            frames.append(pygame.surfarray.pixels3d(screen).swapaxes(0, 1))
+
+        # 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()
+
+        return clip
+
+
+# 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 parabellum import Parabellum, Scenario
+    from jax import random, numpy as jnp
+
+    s = Scenario(jnp.array([[16, 0]]),
+                jnp.array([[0, 32]]) * 8,
+                jnp.zeros((19,), dtype=jnp.uint8),
+                9,
+                 10)
+    env = Parabellum(map_width=32, map_height=32, walls_cause_death=False, scenario=s)
+    rng, key = random.split(random.PRNGKey(0))
+    obs, state = env.reset(key)
+    state_seq = []
+    for step in range(50):
+        rng, key = random.split(rng)
+        key_act = random.split(key, len(env.agents))
+        actions = {
+            agent: jnp.array(1)
+            for i, agent in enumerate(env.agents)
+        }
+        state_seq.append((key, state, actions))
+        rng, key_step = random.split(rng)
+        obs, state, reward, done, infos = env.step(key_step, state, actions)
+
+    vis = Visualizer(env, state_seq)
+    vis.animate()
+
+

parabellum-0.0.0.dist-info/METADATA

@@ -0,0 +1,55 @@
+Metadata-Version: 2.1
+Name: parabellum
+Version: 0.0.0
+Summary: Parabellum environment for parallel warfare simulation
+Home-page: https://github.com/syrkis/parabellum
+License: MIT
+Keywords: warfare,simulation,parallel,environment
+Author: Noah Syrkis
+Author-email: desk@syrkis.com
+Requires-Python: >=3.11,<4.0
+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: darkdetect (>=0.8.0,<0.9.0)
+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: poetry (>=1.8.3,<2.0.0)
+Requires-Dist: pygame (>=2.5.2,<3.0.0)
+Project-URL: Repository, https://github.com/syrkis/parabellum
+Description-Content-Type: text/markdown
+
+# parabellum
+
+Parabellum is an ultra-scalable, high-performance warfare simulation engine.
+It is based on JaxMARL's SMAX environment, but has been heavily modified to
+support a wide range of new features and improvements.
+
+## Installation
+
+Install through PyPI:
+
+```bash
+pip install parabellum
+```
+
+## Usage
+
+```python
+import parabellum as pb
+```
+
+## TODO
+
+- [ ] Parallel pygame vis
+- [ ] Color for health?
+- [ ] Add the ability to see ongoing game.
+- [ ] Bug test friendly fire.
+- [ ] Start sim from arbitrary state.
+- [ ] Save when the episode ends in some state/obs variable
+- [ ] Look for the source of the bug when using more Allies than Enemies
+- [ ] Y inversed axis for parabellum visualization
+- [ ] Units see through obstacles? 
+

parabellum-0.0.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+parabellum/.ipynb_checkpoints/__init__-checkpoint.py,sha256=Yt1RkvkGIJdps0Axpz0ouu-Aaa07032kX04l1l7LXTw,118
+parabellum/.ipynb_checkpoints/env-checkpoint.py,sha256=Z0PD3MJb9Amxl84MMtghTCF92Gr4ln9qSyRx2DSY15Y,11589
+parabellum/.ipynb_checkpoints/vis-checkpoint.py,sha256=7zmFqU99gXSW6ueTeEp3CKMJ9XmrTgJkVEpktdLWd_4,8999
+parabellum/__init__.py,sha256=Yt1RkvkGIJdps0Axpz0ouu-Aaa07032kX04l1l7LXTw,118
+parabellum/env.py,sha256=Z0PD3MJb9Amxl84MMtghTCF92Gr4ln9qSyRx2DSY15Y,11589
+parabellum/vis.py,sha256=7zmFqU99gXSW6ueTeEp3CKMJ9XmrTgJkVEpktdLWd_4,8999
+parabellum-0.0.0.dist-info/METADATA,sha256=cV1VBjjoFLEUmDebGlRpARDdJULpNQ6JspYJ5dqU5ns,1588
+parabellum-0.0.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.0.0.dist-info/RECORD,,

parabellum-0.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.9.0
+Root-Is-Purelib: true
+Tag: py3-none-any