parabellum 0.2.13__py3-none-any.whl → 0.2.14__py3-none-any.whl

parabellum/__init__.py

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

parabellum/env.py

@@ -7,6 +7,7 @@ from jax import random
 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 functools import partial
@@ -16,7 +17,7 @@ from functools import partial
 class Scenario:
     """Parabellum scenario"""
 
-    obstacle_coords: chex.Array
+    obstacle_coords: chex.Array  # TODO: use map instead of obstacles
     obstacle_deltas: chex.Array
 
     unit_types: chex.Array
@@ -30,6 +31,8 @@ class Scenario:
 # default scenario
 scenarios = {
     "default": Scenario(
+        jnp.array([[6, 10], [26, 10]]) * 8,
+        jnp.array([[0, 12], [0, 1]]) * 8,
         jnp.array([[6, 10], [26, 10]]) * 8,
         jnp.array([[0, 12], [0, 1]]) * 8,
         jnp.zeros((19,), dtype=jnp.uint8),
@@ -40,24 +43,17 @@ scenarios = {
 
 
 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)
+    def __init__(self, scenario: Scenario, **kwargs):
+        super(Parabellum, self).__init__(**kwargs)
+        self.obstacle_coords = scenario.obstacle_coords
+        self.obstacle_deltas = scenario.obstacle_deltas
+        self.unit_type_attack_blasts = jnp.zeros((19,), dtype=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
+        self._push_units_away = lambda x: x  # overwrite push units
 
-    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
+    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)
@@ -74,7 +70,7 @@ class Parabellum(SMAX):
             + 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,
@@ -82,7 +78,6 @@ class Parabellum(SMAX):
         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)
@@ -224,19 +219,39 @@ class Parabellum(SMAX):
             perform_agent_action
         )(jnp.arange(self.num_agents), actions, keys)
 
-        # checked that no unit passed through an obstacles 
+        # units push each other
         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))
+        # avoid going into obstacles after being pushed
+
+        bondaries_coords = jnp.array(
+            [[0, 0], [0, 0], [self.map_width, 0], [0, self.map_height]]
+        )
+        bondaries_deltas = jnp.array(
+            [
+                [self.map_width, 0],
+                [0, self.map_height],
+                [0, self.map_height],
+                [self.map_width, 0],
+            ]
+        )
+        obstacle_coords = jnp.concatenate(
+            [self.obstacle_coords, bondaries_coords]
+        )  # add the map boundaries to the obstacles to avoid
+        obstacle_deltas = jnp.concatenate(
+            [self.obstacle_deltas, bondaries_deltas]
+        )  # add the map boundaries to the obstacles to avoid
+        obst_start = obstacle_coords
+        obst_end = obst_start + obstacle_deltas
+
+        def check_obstacles(pos, new_pos, obst_start, obst_end):
+            inters = jnp.any(inter_fn(pos, new_pos, obst_start, obst_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)
-        
+
+        pos = jax.vmap(check_obstacles, in_axes=(0, 0, None, None))(
+            pos, new_pos, obst_start, obst_end
+        )
+
         # Multiple enemies can attack the same unit.
         # We have `(health_diff, attacked_idx)` pairs.
         # `jax.lax.scatter_add` aggregates these exactly
@@ -278,19 +293,23 @@ class Parabellum(SMAX):
         )
         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)
+    n_envs = 4
+    kwargs = dict(map_width=64, map_height=64)
+    env = Parabellum(scenarios["default"], **kwargs)
+    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 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))
-
 
+    for i in range(10):
+        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)

parabellum/map.py

@@ -0,0 +1,16 @@
+# map.py
+#   parabellum map functions
+# by: Noah Syrkis
+
+# imports
+import jax.numpy as jnp
+import jax
+
+
+# functions
+def map_fn(width, height, obst_coord, obst_delta):
+    """Create a map from the given width, height, and obstacle coordinates and deltas."""
+    m = jnp.zeros((width, height))
+    for (x, y), (dx, dy) in zip(obst_coord, obst_delta):
+        m = m.at[x : x + dx, y : y + dy].set(1)
+    return m

parabellum/run.py

@@ -0,0 +1,125 @@
+# run.py
+#   parabellum run game live
+# by: Noah Syrkis
+
+# Noah Syrkis
+import pygame
+from jax import random
+from functools import partial
+import darkdetect
+import jax.numpy as jnp
+from chex import dataclass
+import jaxmarl
+from typing import Tuple, List, Dict, Optional
+import parabellum as pb
+
+
+# constants
+fg = (255, 255, 255) if darkdetect.isDark() else (0, 0, 0)
+bg = (0, 0, 0) if darkdetect.isDark() else (255, 255, 255)
+
+
+# types
+State = jaxmarl.environments.smax.smax_env.State
+Obs = Reward = Done = Action = Dict[str, jnp.ndarray]
+StateSeq = List[Tuple[jnp.ndarray, State, Action]]
+
+
+@dataclass
+class Control:
+    running: bool = True
+    paused: bool = False
+    click: Optional[Tuple[int, int]] = None
+
+
+@dataclass
+class Game:
+    clock: pygame.time.Clock
+    state: State
+    obs: Dict
+    state_seq: StateSeq
+    control: Control
+    env: pb.Parabellum
+    rng: random.PRNGKey
+
+
+def handle_event(event, control_state):
+    """Handle pygame events."""
+    if event.type == pygame.QUIT:
+        control_state.running = False
+    if event.type == pygame.MOUSEBUTTONDOWN:
+        pos = pygame.mouse.get_pos()
+        control_state.click = pos
+    if event.type == pygame.MOUSEBUTTONUP:
+        control_state.click = None
+    if event.type == pygame.KEYDOWN:  # any key press pauses
+        control_state.paused = not control_state.paused
+    return control_state
+
+
+def control_fn(game):
+    """Handle pygame events."""
+    for event in pygame.event.get():
+        game.control = handle_event(event, game.control)
+    return game
+
+
+def render_fn(screen, game):
+    """Render the game."""
+    if len(game.state_seq) < 3:
+        return game
+    for rng, state, action in env.expand_state_seq(game.state_seq[-2:])[-8:]:
+        screen.fill(bg)
+        if game.control.click is not None:
+            pygame.draw.circle(screen, "red", game.control.click, 10)
+        unit_positions = state.unit_positions
+        for pos in unit_positions:
+            pos = (pos / env.map_width * 800).tolist()
+            pygame.draw.circle(screen, fg, pos, 5)
+        pygame.display.flip()
+        game.clock.tick(24)  # limits FPS to 24
+    return game
+
+
+def step_fn(game):
+    """Step in parabellum."""
+    rng, act_rng, step_key = random.split(game.rng, 3)
+    act_key = random.split(act_rng, env.num_agents)
+    action = {
+        a: env.action_space(a).sample(act_key[i]) for i, a in enumerate(env.agents)
+    }
+    state_seq_entry = (step_key, game.state, action)
+    # append state_seq_entry to state_seq
+    game.state_seq.append(state_seq_entry)
+    obs, state, reward, done, info = env.step(step_key, game.state, action)
+    game.state = state
+    game.obs = obs
+    game.rng = rng
+    return game
+
+
+# state
+if __name__ == "__main__":
+    env = pb.Parabellum(pb.scenarios["default"])
+    pygame.init()
+    screen = pygame.display.set_mode((1000, 1000))
+    render = partial(render_fn, screen)
+    rng, key = random.split(random.PRNGKey(0))
+    obs, state = env.reset(key)
+    kwargs = dict(
+        control=Control(),
+        env=env,
+        rng=rng,
+        state_seq=[],  # [(key, state, action)]
+        clock=pygame.time.Clock(),
+        state=state,
+        obs=obs,
+    )
+    game = Game(**kwargs)
+
+    while game.control.running:
+        game = control_fn(game)
+        game = game if game.control.paused else step_fn(game)
+        game = game if game.control.paused else render(game)
+
+    pygame.quit()

parabellum/vis.py

@@ -4,9 +4,11 @@ from tqdm import tqdm
 import jax.numpy as jnp
 import jax
 from jax import vmap
+from jax import tree_util
 from functools import partial
 import darkdetect
 import pygame
+import os
 from moviepy.editor import ImageSequenceClip
 from typing import Optional
 from jaxmarl.environments.multi_agent_env import MultiAgentEnv
@@ -20,6 +22,14 @@ from collections import defaultdict
 action_to_symbol = {0: "↑", 1: "→", 2: "↓", 3: "←", 4: "Ø"}
 
 
+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))
+
+
 class Visualizer(SMAXVisualizer):
     def __init__(self, env: MultiAgentEnv, state_seq, reward_seq=None):
         super().__init__(env, state_seq, reward_seq)
@@ -30,7 +40,54 @@ class Visualizer(SMAXVisualizer):
         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)
+        # self.bullet_seq = vmap(partial(bullet_fn, self.env))(self.state_seq)
+
+    def animate(self, save_fname: str = "output/parabellum.mp4"):
+        multi_dim = self.state_seq[0][1].unit_positions.ndim > 1
+        if multi_dim:
+            n_envs = self.state_seq[0][1].unit_positions.shape[0]
+            if not self.have_expanded:
+                state_seqs = vmap(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 = 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
+        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
+
+            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
+            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
 
     def render_agents(self, screen, state):
         time_tuple = zip(
@@ -100,39 +157,6 @@ class Visualizer(SMAXVisualizer):
             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
@@ -201,30 +225,30 @@ def bullet_fn(env, states):
 
 # test the visualizer
 if __name__ == "__main__":
-    from parabellum import Parabellum, Scenario
     from jax import random, numpy as jnp
+    from parabellum import Parabellum, scenarios
+
+    # small_multiples()  # testing small multiples (not working yet)
+    # exit()
 
-    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)
+    n_envs = 100
+    kwargs = dict(map_width=64, map_height=64)
+    env = Parabellum(scenarios["default"], **kwargs)
+    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 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)
+
+    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)
         }
-        state_seq.append((key, state, actions))
-        rng, key_step = random.split(rng)
-        obs, state, reward, done, infos = env.step(key_step, state, actions)
+        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()
-
-

parabellum-0.2.14.dist-info/METADATA

@@ -0,0 +1,104 @@
+Metadata-Version: 2.1
+Name: parabellum
+Version: 0.2.14
+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: 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: poetry (>=1.8.3,<2.0.0)
+Requires-Dist: pygame (>=2.5.2,<3.0.0)
+Requires-Dist: tqdm (>=4.66.4,<5.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
+
+Parabellum is written in Python 3.11 and can be installed using pip:
+
+```bash
+pip install parabellum
+```
+
+## Usage
+
+Parabellum is designed to be used in conjunction with JAX, a high-performance
+numerical computing library. Here is a simple example of how to use Parabellum
+to simulate a game with 10 agents and 10 enemies, each taking random actions:
+
+```python
+import parabellum as pb
+from jax import random
+
+# define the scenario
+kwargs = dict(obstacle_coords=[(7, 7)], obstacle_deltas=[(10, 0)])
+scenario = pb.Scenario(**kwargs)  # <- Scenario is an important part of parabellum
+
+# create the environment
+kwargs = dict(map_width=256, map_height=256, num_agents=10, num_enemies=10)
+env = pb.Parabellum(**kwargs)  # <- Parabellum is the central class of parabellum
+
+# initiate stochasticity
+rng = random.PRNGKey(0)
+rng, key = random.split(rng)
+
+# initialize the environment state
+obs, state = env.reset(key)
+state_sequence = []
+
+for _ in range(1000):
+
+    # manage stochasticity
+    rng, rng_act, key_step = random.split(key)
+    key_act = random.split(rng_act, len(env.agents))
+
+    # sample actions and append to state sequence
+    act = {a: env.action_space(a).sample(k)
+        for a, k in zip(env.agents, key_act)}
+
+    # step the environment
+    state_sequence.append((key_act, state, act))
+    obs, state, reward, done, info = env.step(key_step, act, state)
+
+
+# save visualization of the state sequence
+vis = pb.Visualizer(env, state_sequence)  # <- Visualizer is a nice to have class
+vis.animate()
+```
+
+
+## Features
+
+- Obstacles — can be inserted in
+
+## TODO
+
+- [x] Parallel pygame vis
+    - [ ] Parallel bullet renderings
+    - [ ] Combine parallell plots into one (maybe out of parabellum scope)
+- [ ] Color for health?
+- [ ] Add the ability to see ongoing game.
+- [ ] Bug test friendly fire.
+- [x] 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.2.14.dist-info/RECORD

@@ -0,0 +1,8 @@
+parabellum/__init__.py,sha256=TjZVlHZdi7CEm0gjagm3j6epPZxjR6C9G3CYtX7d-2o,142
+parabellum/env.py,sha256=rCn6iPLeFpqitncD9nEc0KA6N9JCMmiSyP9u2meOJxk,12325
+parabellum/map.py,sha256=SQeNl1kkGsnnqYoo-60zJNv36fD-8VSKasiS1_WARao,410
+parabellum/run.py,sha256=lVNBsMc8HY4Tqdjs_1MXGBvIzuN05brbRiqp0xlRc6c,3410
+parabellum/vis.py,sha256=u7ifxWzHf96WgLTz_hw0ijy6-7wePd7lf0p-yD-NCQY,10212
+parabellum-0.2.14.dist-info/METADATA,sha256=wEiXzwPfnigG5ZSANPFwGEjLCDU5D0c7qbpvEi6Gbm8,3223
+parabellum-0.2.14.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.2.14.dist-info/RECORD,,

parabellum/.ipynb_checkpoints/__init__-checkpoint.py

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

parabellum/.ipynb_checkpoints/env-checkpoint.py

@@ -1,296 +0,0 @@
-"""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

@@ -1,230 +0,0 @@
-"""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.2.13.dist-info/METADATA

@@ -1,56 +0,0 @@
-Metadata-Version: 2.1
-Name: parabellum
-Version: 0.2.13
-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)
-Requires-Dist: tqdm (>=4.66.4,<5.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.2.13.dist-info/RECORD

@@ -1,9 +0,0 @@
-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.2.13.dist-info/METADATA,sha256=jIs4QuEQ7Act04HK2cwKw4PaXOi9Y01ucQ2a4np4-KU,1627
-parabellum-0.2.13.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-parabellum-0.2.13.dist-info/RECORD,,