Running "humanoid_joystick.ipynb" reports an error

[fushengcheng]

When training in the G1JoystickFlatTerrain environment, an error is reported as shown in the figure. How can I fix it?

[kevinzakka]

I need more information. What GPU version? How many GPUs? Do you have a heterogeneous GPU setup (aka different GPUs on the same machine)? What JAX version?

[fushengcheng]

I need more information. What GPU version? How many GPUs? Do you have a heterogeneous GPU setup (aka different GPUs on the same machine)? What JAX version?

Thank you for your reply. The nvidia dirver version is 535.54.03 and CUDA version is 12.2. I have 2 homogeneous GPUs on the same machine. JAX version is 0.5.3.
I've also tried it on a single GPU machine and I get this error as well

[Specialized123]

I have this problem too,I use nvidia 4090 24G, NVIDIA-SMI:570.124.04,Driver Version: 570.124.04 ,CUDA Version: 12.8; JAX version : 0.5.3

[fushengcheng]

I have this problem too,I use nvidia 4090 24G, NVIDIA-SMI:570.124.04,Driver Version: 570.124.04 ,CUDA Version: 12.8; JAX version : 0.5.3

You can try using version 0.4.35 of JAX, which can solve this problem in some environments. However, when I attempted to incorporate the corresponding training environment of Go2 by referring to the joystick and getup environment of Go1, the getup environment still threw this error under the old version of JAX.

[cvoelcker]

Hi, I also (finally) tested this, and got the same results. Changing the jax versions fixes nan issues in some envs, others, like humanoid, still produce nan. Tested with 0.5.3 and 0.4.35

[42jaylonw]

Hi, I also (finally) tested this, and got the same results. Changing the jax versions fixes nan issues in some envs, others, like humanoid, still produce nan. Tested with 0.5.3 and 0.4.35

I got same problem when randomizing body_mass in G1 environment, but these version solved it !

python==3.13.2
jax==0.4.38
jax-cuda12-pjrt==0.4.38
jax-cuda12-plugin==0.4.38
jaxlib==0.4.38
jaxopt==0.8.5
flax==0.10.4
brax==0.12.3

[Hojunian]

I got the same problem when using rtx 4090, and solved the issue without changing JAX version.

I just added jax.config.update('jax_default_matmul_precision', 'highest'), following google-deepmind/mujoco#1484

[Update (May 23)]
I found that DM Control Suite Humanoid environments produce nan even after changing the matmul precision. Changing JAX version following the solution of @42jaylonw seems to resolve this issue. I will report complete results after running several ppo training.