fix: fix cond_bug and changed int type

RobertTLange · Sep 1, 2023 · edaf650 · edaf650
1 parent ad9abac
commit edaf650
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gymnax/environments/misc/rooms.py b/gymnax/environments/misc/rooms.py
@@ -111,7 +111,7 @@ def step_env(
         # Caculate the disocunted return
         discounted_reward = normal_reward - 0.9 *( state.time / params.max_steps_in_episode)
 
-        reward = jax.lax.cond(params.discounted_reward, lambda: discounted_reward, lambda: int(normal_reward))
+        reward = jax.lax.cond(params.discounted_reward, lambda: discounted_reward, lambda: jnp.int32(normal_reward))
         # Update state dict and evaluate termination conditions
         state = EnvState(new_pos, state.goal, state.time + 1)
         done = self.is_terminal(state, params)