parabellum 0.2.25__py3-none-any.whl → 0.2.26__py3-none-any.whl

parabellum/env.py

@@ -60,7 +60,7 @@ def make_scenario(place, terrain_raster, unit_starting_sectors, allies_type, n_a
     else:
         assert(len(allies_type) == n_allies)
         allies = allies_type
-        
+
     if type(enemies_type) == int:
         enemies = [enemies_type] * n_enemies
     else:
@@ -149,9 +149,14 @@ class Environment(SMAX):
         state = self._push_units_away(state)  # type: ignore
         obs = self.get_obs(state)
         world_state = self.get_world_state(state)
-        obs["world_state"] = jax.lax.stop_gradient(world_state)
+        # obs["world_state"] = jax.lax.stop_gradient(world_state)
         return obs, state
 
+    def step_env(self, rng, state: State, action: Array):
+        obs, state, rewards, dones, infos = super().step_env(rng, state, action)
+        # delete world_state from obs
+        obs.pop("world_state")
+        return obs, state, rewards, dones, infos
 
     def _our_push_units_away(
         self, pos, unit_types, firmness: float = 1.0
@@ -198,7 +203,7 @@ class Environment(SMAX):
             raster = jnp.where(jnp.arange(raster.shape[1]) >= minimum[1], raster.T, 0).T
             raster = jnp.where(jnp.arange(raster.shape[1]) <= maximum[1], raster.T, 0).T
             return jnp.any(raster)
-        
+
         def update_position(idx, vec):
             # Compute the movements slightly strangely.
             # The velocities below are for diagonal directions
@@ -337,7 +342,7 @@ class Environment(SMAX):
         pos = jax.vmap(jnp.where)(clash, pos, new_pos)
         # avoid going out of bounds
         pos = jnp.maximum(jnp.minimum(pos, jnp.array([self.map_width, self.map_height])),jnp.zeros((2,)),)
-        
+
         # Multiple enemies can attack the same unit.
         # We have `(health_diff, attacked_idx)` pairs.
         # `jax.lax.scatter_add` aggregates these exactly

{parabellum-0.2.25.dist-info → parabellum-0.2.26.dist-info}/METADATA

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

{parabellum-0.2.25.dist-info → parabellum-0.2.26.dist-info}/RECORD

@@ -1,10 +1,10 @@
 parabellum/__init__.py,sha256=-5cWXJkHnfH_CbhTEall8Wak8McAFXZHP1L8Fu7Uo5k,373
 parabellum/aid.py,sha256=HWST27inTFXcp8b11izJF0U7N7DZnRTIS3n1Qfa-Ko4,106
-parabellum/env.py,sha256=H1VwxHj8KqbjBJ8b7NMxkAg3Q0qwVzGulrigc26Tzkc,16663
+parabellum/env.py,sha256=H1YGAtUYNJd8OHnZ3sOEXbag5L0WjtJHBGL8ymGPvoE,16898
 parabellum/gun.py,sha256=nvsJdcZ2Qd6lbPlAgsUiaLhstTi1UdLQ8kOnbCenucY,2618
 parabellum/map.py,sha256=UwMqwySasX5oLw9v5YMJARAPwvQThLTRW36NpbwvBC8,3564
 parabellum/run.py,sha256=EO_F7VPwayatpSHrcbSahtinsV4QObhcx0jo-4KZO1E,3472
 parabellum/vis.py,sha256=q7_OIMjzt-7nBOojVVW7Wiiq9ojsjaltIsH6eOOxPKk,6116
-parabellum-0.2.25.dist-info/METADATA,sha256=UCuRLYhSUxnebs5pRs1FaWK67PUJw0kgsJ2UaaxNM9Q,2671
-parabellum-0.2.25.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-parabellum-0.2.25.dist-info/RECORD,,
+parabellum-0.2.26.dist-info/METADATA,sha256=AJwdmHRRPG2MosgufeWqQH7LsDRtTvFFGMh1azey9zA,2671
+parabellum-0.2.26.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+parabellum-0.2.26.dist-info/RECORD,,