warp backend causes physics to go to nan

[aaprasad]

I've been using mujoco-playground to train a C.elegans biomechanical model to do imitation learning using the mjx backend and it has worked fine. When I try switching to warp with identitical configurations otherwise, the physics becomes unstable and goes to nan. I'm not sure what could be causing this. These are the library versions I'm working with:

jax                                     0.7.2               pypi_0                 pypi
jax-cuda12-pjrt                         0.7.2               pypi_0                 pypi
jax-cuda12-plugin                       0.7.2               pypi_0                 pypi
jaxlib                                  0.7.2               pypi_0                 pypi
jaxopt                                  0.8.5               pypi_0                 pypi
mujoco                                  3.4.0            pypi_0                 pypi
mujoco-mjx                              3.4.0            pypi_0                 pypi
mujoco-warp                             0.0.1            pypi_0                 pypi
playground                              0.0.5            pypi_0                 pypi
warp-lang                               1.10.0           pypi_0                 pypi

I'm on ubuntu and using a 5090 gpu with cuda 12.8. I've attached a minimal example to reproduce the bug by just stepping through a basic environment:

warp_debug.zip

[thowell]

@aaprasad thank you for opening an issue and providing a repro!

please try installing mujoco warp from source and then simulating the scene in the mujoco warp viewer

mjwarp-viewer celegans_reduce_contacts.xml --nconmax=256 --njmax=512

does the scene simulate as expected in the viewer? thanks!

[aaprasad]

Hi @thowell, thanks for the quick reply! I extracted the spec created by the environment in the notebook and loaded it into the viewer but it was all nans. However, I was able to load it in the cpu mujoco viewer and the mjx viewer with jax backend... Below is a screenshot of the warp viewer vs a screenshot of the cpu mujoco viewer.

spec extracted from env: test_celegans_spec.xml

.

Interestingly, when I try loading just the body model without the floor (ie celegans_reduce_contacts.xml) it loads fine

Screencast.from.2025-12-17.15-57-22.1.webm