add rng to env step

syrkis · Aug 7, 2024 · 74a28f5 · 74a28f5
1 parent 4846247
commit 74a28f5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/parabellum/env.py b/parabellum/env.py
@@ -129,8 +129,8 @@ def reset(self, rng: chex.PRNGKey) -> Tuple[Dict[str, chex.Array], State]:
         # obs["world_state"] = jax.lax.stop_gradient(world_state)
         return obs, state
 
-    def step_env(self, state: State, action: Array):
-        obs, state, rewards, dones, infos = super().step_env(state, action)
+    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