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

parabellum/env.py

@@ -17,6 +17,8 @@ from functools import partial
 class Scenario:
     """Parabellum scenario"""
 
+    terrain_raster: chex.Array
+
     obstacle_coords: chex.Array  # TODO: use map instead of obstacles
     obstacle_deltas: chex.Array
 
@@ -31,10 +33,9 @@ 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.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,
@@ -44,13 +45,74 @@ scenarios = {
 
 class Parabellum(SMAX):
     def __init__(self, scenario: Scenario, **kwargs):
-        super(Parabellum, self).__init__(**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
         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
@@ -79,13 +141,24 @@ class Parabellum(SMAX):
         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
@@ -109,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)
 
             #######################################################################
             #######################################################################
@@ -245,8 +320,10 @@ class Parabellum(SMAX):
         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)
+            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

parabellum/run.py

@@ -123,3 +123,5 @@ if __name__ == "__main__":
         game = game if game.control.paused else render(game)
 
     pygame.quit()
+
+

parabellum/vis.py

@@ -7,6 +7,7 @@ from jax import vmap
 from jax import tree_util
 from functools import partial
 import darkdetect
+import numpy as np
 import pygame
 import os
 from moviepy.editor import ImageSequenceClip
@@ -47,7 +48,7 @@ class Visualizer(SMAXVisualizer):
         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)
+                state_seqs = vmap(self.env.expand_state_seq)(self.state_seq)
                 self.have_expanded = True
             for i in range(n_envs):
                 state_seq = jax.tree_map(lambda x: x[i], state_seqs)
@@ -62,12 +63,19 @@ class Visualizer(SMAXVisualizer):
     def animate_one(self, state_seq, action_seq, save_fname):
         frames = []  # frames for the video
         pygame.init()  # initialize pygame
+        terrain = np.array(self.env.terrain_raster)
+        rgb_array = np.zeros((terrain.shape[0], terrain.shape[1], 3), dtype=np.uint8)
+        rgb_array[terrain == 1] = self.fg
+        mask_surface = pygame.surfarray.make_surface(rgb_array)
+        mask_surface = pygame.transform.scale(mask_surface, (self.s, self.s))
+
         for idx, (_, state, _) in tqdm(enumerate(state_seq), total=len(self.state_seq)):
             action = action_seq[idx // self.env.world_steps_per_env_step]
             screen = pygame.Surface(
                 (self.s, self.s), pygame.HWSURFACE | pygame.DOUBLEBUF
             )
             screen.fill(self.bg)  # fill the screen with the background color
+            screen.blit(mask_surface, (0, 0))
 
             self.render_agents(screen, state)  # render the agents
             self.render_action(screen, action)
@@ -80,7 +88,6 @@ class Visualizer(SMAXVisualizer):
 
             # 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)
@@ -99,7 +106,6 @@ class Visualizer(SMAXVisualizer):
         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]
@@ -231,9 +237,8 @@ if __name__ == "__main__":
     # small_multiples()  # testing small multiples (not working yet)
     # exit()
 
-    n_envs = 100
-    kwargs = dict(map_width=64, map_height=64)
-    env = Parabellum(scenarios["default"], **kwargs)
+    n_envs = 2
+    env = Parabellum(scenarios["default"])
     rng, reset_rng = random.split(random.PRNGKey(0))
     reset_key = random.split(reset_rng, n_envs)
     obs, state = vmap(env.reset)(reset_key)
@@ -243,7 +248,7 @@ if __name__ == "__main__":
         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])
+            a: jnp.ones_like(vmap(env.action_space(a).sample)(act_key[i]))
             for i, a in enumerate(env.agents)
         }
         step_key = random.split(step_rng, n_envs)

{parabellum-0.2.14.dist-info → parabellum-0.2.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: parabellum
-Version: 0.2.14
+Version: 0.2.16
 Summary: Parabellum environment for parallel warfare simulation
 Home-page: https://github.com/syrkis/parabellum
 License: MIT

parabellum-0.2.16.dist-info/RECORD

@@ -0,0 +1,8 @@
+parabellum/__init__.py,sha256=TjZVlHZdi7CEm0gjagm3j6epPZxjR6C9G3CYtX7d-2o,142
+parabellum/env.py,sha256=d6agGy-kTRIg_r0QKCL_7iztzwhaTfsb4yhtUQfdgx0,16024
+parabellum/map.py,sha256=SQeNl1kkGsnnqYoo-60zJNv36fD-8VSKasiS1_WARao,410
+parabellum/run.py,sha256=0IWpqcQ_qfFeElbupF5vOs_CByFfpXYuGGUHYuurFM4,3412
+parabellum/vis.py,sha256=euT7VNPpKW9h0bjXwtYBa4MJRXuELfH3JnUm5ulr3s0,10559
+parabellum-0.2.16.dist-info/METADATA,sha256=eXEfS4FXFp4Xrp4g3hrKnvh-fIHzmHcWlnZrIRjdF4k,3223
+parabellum-0.2.16.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.2.16.dist-info/RECORD,,

parabellum-0.2.14.dist-info/RECORD

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