Error running the Asterix environment

[ID-Akash]

Hello,

I am trying to train a RL agent in JAX for Asterix-Minatar, however, I get the following error:

File "/opt/miniconda3/envs/minasa/lib/python3.11/site-packages/gymnax/wrappers/gym.py", line 70, in step o, self.env_state, r, d, info = self._env.step( ^^^^^^^^^^^^^^^ File "/opt/miniconda3/envs/minasa/lib/python3.11/site-packages/gymnax/environments/environment.py", line 45, in step obs_st, state_st, reward, done, info = self.step_env(key, state, action, params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/opt/miniconda3/envs/minasa/lib/python3.11/site-packages/gymnax/environments/minatar/asterix.py", line 108, in step_env state, reward, done = step_entities(state) ^^^^^^^^^^^^^^^^^^^^ File "/opt/miniconda3/envs/minasa/lib/python3.11/site-packages/gymnax/environments/minatar/asterix.py", line 335, in step_entities bool(done > 0), ^^^^^^^^^^^^^^ jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].. The error occurred while tracing the function step at /opt/miniconda3/envs/minasa/lib/python3.11/site-packages/gymnax/environments/environment.py:32 for jit. This concrete value was not available in Python because it depends on the values of the arguments self, key.player_x, key.player_y, key.spawn_timer, key.move_timer, key.entities, and state. See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

The same code works fine with SpaceInvaders and Breakout environments. I think that because the done flag depends on the state, jax is unable to jit the step function.

I am using python 3.11 and jax 0.4.23.

[JackZhangY]

Refer to the code in the 'main' branch to fix some small debugs that exist in version v0.0.8,
e.g., line 335 in 'asterix.py':
bool(done > 0) --> jnp.bool_(done > 0)