Merge pull request #74 from EdanToledo/main

fix: non-jittable code for minatar freeway and asterix
RobertTLange · May 1, 2024 · aef77d5 · aef77d5
2 parents 88fef69 + 0096e0e
commit aef77d5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/gymnax/environments/minatar/asterix.py b/gymnax/environments/minatar/asterix.py
@@ -332,7 +332,7 @@ def step_entities(
     return (
         state.replace(entities=entities, move_timer=move_timer),
         reward,
-        bool(done > 0),
+        jnp.bool_(done > 0),
     )
 
 

diff --git a/gymnax/environments/minatar/freeway.py b/gymnax/environments/minatar/freeway.py
@@ -199,7 +199,7 @@ def step_agent(
     win_cond = pos == 0
     reward = win_cond * 1.0
     pos = jax.lax.select(win_cond, 9, pos)
-    return state.replace(pos=pos, move_timer=move_timer), reward, win_cond.item()
+    return state.replace(pos=pos, move_timer=move_timer), reward, win_cond
 
 
 def step_cars(state: EnvState) -> EnvState: