google-deepmind / google-deepmind/mujoco_playground
Training performance collapse/nan in HumanoidRun/Walk/Stand
- Dominant language
- Python
- Stars
- 2.2k
- Forks
- 359
- Avg merge
- 7d 3h
- Merged PRs (30d)
- 1
Description
Following up from #92 as @btaba suggested this is a separate issue.
The DMC Humanoid envs (run/walk/stand) result in nans in training after a while. Details from the other thread.
Info from the other thread
Hi, after doing extensive tests yesterday it seems like neither fix (concerning symmetric Cholesky decomposition) reliably solves the issue on our server. It does reduce the instances of it occurring though.
Specs are an single L40s GPU, Driver Version: 570.148.08, CUDA Version: 12.8, jax[cuda12]==0.6.2 (with manual fix), mujoco==3.3.3 (with manual fix), playground==0.0.5. I am running HumanoidRun, HumanoidWalk, and HumanoidStand from the DMC implementation shipped with mujoco_playground. I tested both fixes together and either in isolation. If necessary I can try to get some stats on how many runs fail with either fix. Somehow, HumanoidWalk seems to be the most affected.
Full list of jax related libraries installed
```
distrax==0.1.5
flax==0.10.6
gymnax==0.0.9
jax==0.6.2
jax-cuda12-pjrt==0.6.2
jax-cuda12-plugin==0.6.2
jaxlib==0.6.2
jaxopt==0.8.5
mujoco==3.3.3
playground==0.0.5
```
- Training for a long time (several million timesteps). Failures happen both with very good and bad policies, no clear difference. I am not running PPO but a custom algorithm, but it seems extremely stable in all other envs (brax and playground) so I am willing to bet that it is not an algorithm problem.
- Training ends with nans. I am specifically guarding against feeding nan actions to the environment (jax.nan_to_num), so as far as I can tell the source of the nans is the environment.
- I tried roughly 30 runs per fix combination, all of them end in a failure at some point over a maximum of 200_000 steps with 1024 envs in parallel (so 200_000_000 timesteps in total)
- It is as far as I can tell only HumanoidRun/Walk/Stand. I tested all playground DMC envs, and G1 and T1 Joystick Humanoid rough and flat terrain. Other DMC envs and the G1 and T1 environments seem fine.
- jax==0.4.38 showed the same issues. I was unable to test jax==0.4.35 due to dependency conflicts with other libraries.
- I did set the precision to highest. I can repeat the test on one of our 3090 GPUs.
I verified that this error is reproducible on an L40s. Test on the 4090 are running.
Contributor guide
Assessment
This issue has not been assessed yet.