parabellum 0.2.13__tar.gz → 0.2.15__tar.gz

parabellum-0.2.15/PKG-INFO

@@ -0,0 +1,104 @@
+Metadata-Version: 2.1
+Name: parabellum
+Version: 0.2.15
+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.15/README.md

@@ -0,0 +1,78 @@
+# 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.15/parabellum/__init__.py

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

{parabellum-0.2.13 → parabellum-0.2.15}/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,9 @@ from functools import partial
 class Scenario:
     """Parabellum scenario"""
 
-    obstacle_coords: chex.Array
+    terrain_raster: chex.Array
+
+    obstacle_coords: chex.Array  # TODO: use map instead of obstacles
     obstacle_deltas: chex.Array
 
     unit_types: chex.Array
@@ -30,8 +33,9 @@ class Scenario:
 # default scenario
 scenarios = {
     "default": Scenario(
-        jnp.array([[6, 10], [26, 10]]) * 8,
-        jnp.array([[0, 12], [0, 1]]) * 8,
+        jnp.eye(128, dtype=jnp.uint8),
+        jnp.array([[80, 0], [16, 12]]),
+        jnp.array([[0, 80], [0, 20]]),
         jnp.zeros((19,), dtype=jnp.uint8),
         9,
         10,
@@ -40,24 +44,78 @@ 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):
+        map_height, map_width = scenario.terrain_raster.shape
+        args = dict(scenario=scenario, map_height=map_height, map_width=map_width)
+        super(Parabellum, self).__init__(**args, **kwargs)
+        self.terrain_raster = scenario.terrain_raster
+        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
-
-    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
+        self._push_units_away = lambda x: x  # overwrite push units
+
+    @partial(jax.jit, static_argnums=(0,))
+    def reset(self, key: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:
+        """Environment-specific reset."""
+        key, team_0_key, team_1_key = jax.random.split(key, num=3)
+        team_0_start = jnp.stack(
+            [jnp.array([self.map_width / 4, self.map_height / 2])] * self.num_allies
+        )
+        team_0_start_noise = jax.random.uniform(
+            team_0_key, shape=(self.num_allies, 2), minval=-2, maxval=2
+        )
+        team_0_start = team_0_start + team_0_start_noise
+        team_1_start = jnp.stack(
+            [jnp.array([self.map_width / 4 * 3, self.map_height / 2])]
+            * self.num_enemies
+        )
+        team_1_start_noise = jax.random.uniform(
+            team_1_key, shape=(self.num_enemies, 2), minval=-2, maxval=2
+        )
+        team_1_start = team_1_start + team_1_start_noise
+        unit_positions = jnp.concatenate([team_0_start, team_1_start])
+        key, pos_key = jax.random.split(key)
+        generated_unit_positions = self.position_generator.generate(pos_key)
+        unit_positions = jax.lax.select(
+            self.smacv2_position_generation, generated_unit_positions, unit_positions
+        )
+        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 = (
+            jnp.zeros((self.num_agents,), dtype=jnp.uint8)
+            if self.scenario is None
+            else self.scenario
+        )
+        key, unit_type_key = jax.random.split(key)
+        generated_unit_types = self.unit_type_generator.generate(unit_type_key)
+        unit_types = jax.lax.select(
+            self.smacv2_unit_type_generation, generated_unit_types, unit_types
+        )
+        unit_health = self.unit_type_health[unit_types]
+        state = State(
+            unit_positions=unit_positions,
+            unit_alive=jnp.ones((self.num_agents,), dtype=jnp.bool_),
+            unit_teams=unit_teams,
+            unit_health=unit_health,
+            unit_types=unit_types,
+            prev_movement_actions=jnp.zeros((self.num_agents, 2)),
+            prev_attack_actions=jnp.zeros((self.num_agents,), dtype=jnp.int32),
+            time=0,
+            terminal=False,
+            unit_weapon_cooldowns=unit_weapon_cooldowns,
+        )
+        state = self._push_units_away(state)
+        obs = self.get_obs(state)
+        world_state = self.get_world_state(state)
+        obs["world_state"] = jax.lax.stop_gradient(world_state)
+        return obs, 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)
@@ -74,7 +132,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,15 +140,25 @@ 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):
+        def intersect_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 raster_crossing(pos, new_pos):
+            pos, new_pos = pos.astype(jnp.int32), new_pos.astype(jnp.int32)
+            raster = self.terrain_raster
+            axis = jnp.argmax(jnp.abs(new_pos - pos), axis=-1)
+            minimum = jnp.minimum(pos[axis], new_pos[axis]).squeeze()
+            maximum = jnp.maximum(pos[axis], new_pos[axis]).squeeze()
+            segment = jnp.where(axis == 0, raster[pos[1]], raster.T[pos[0]])
+            segment = jnp.where(jnp.arange(segment.shape[0]) >= minimum, segment, 0)
+            segment = jnp.where(jnp.arange(segment.shape[0]) <= maximum, segment, 0)
+            return jnp.any(segment)
+
         def update_position(idx, vec):
             # Compute the movements slightly strangely.
             # The velocities below are for diagonal directions
@@ -114,8 +182,10 @@ class Parabellum(SMAX):
             ############################################ 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)
+            inters = jnp.any(intersect_fn(pos, new_pos, obs, obs_end))
+            rastersects = raster_crossing(pos, new_pos)
+            flag = jnp.logical_or(inters, rastersects)
+            new_pos = jnp.where(flag, pos, new_pos)
 
             #######################################################################
             #######################################################################
@@ -224,19 +294,41 @@ 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))
-            return jnp.where(inters, pos, new_pos)
-        
-        pos = jax.vmap(check_obstacles, in_axes=(0,0,None,None))(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):
+            intersects = jnp.any(intersect_fn(pos, new_pos, obst_start, obst_end))
+            rastersect = raster_crossing(pos, new_pos)
+            flag = jnp.logical_or(intersects, rastersect)
+            return jnp.where(flag, pos, new_pos)
+
+        pos = jax.vmap(check_obstacles, in_axes=(0, 0, None, None))(
+            pos, new_pos, obst_start, obst_end
+        )
+
         # Multiple enemies can attack the same unit.
         # We have `(health_diff, attacked_idx)` pairs.
         # `jax.lax.scatter_add` aggregates these exactly
@@ -278,19 +370,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-0.2.15/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-0.2.15/parabellum/run.py

@@ -0,0 +1,127 @@
+# 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()
+
+