add rng to env step

parabellum/env.py

@@ -126,14 +126,9 @@ def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:
         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