parabellum 0.2.21__tar.gz → 0.2.23__tar.gz

{parabellum-0.2.21 → parabellum-0.2.23}/PKG-INFO

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

{parabellum-0.2.21 → parabellum-0.2.23}/parabellum/env.py

@@ -97,6 +97,7 @@ class Environment(SMAX):
         self.top_sector, self.top_sector_offset = sector_fn(self.terrain_raster, 0)
         self.low_sector, self.low_sector_offset = sector_fn(self.terrain_raster, 24)
 
+
     @partial(jax.jit, static_argnums=(0,))
     def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:
         """Environment-specific reset."""
@@ -125,9 +126,15 @@ 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
     ):  # copy of SMAX._push_units_away but used without state and called inside _world_step to allow more obstacles constraints

{parabellum-0.2.21 → parabellum-0.2.23}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "parabellum"
-version = "v0.2.21"
+version = "v0.2.23"
 description = "Parabellum environment for parallel warfare simulation"
 authors = ["Noah Syrkis <desk@syrkis.com>"]
 readme = "README.md"