google-deepmind / google-deepmind/mujoco
MJX with a minimal quadruped-leg model: step produces NaN while MuJoCo native is stable
- Dominant language
- C++
- Stars
- 15.2k
- Forks
- 1.8k
- Avg merge
- 10d 16h
- Merged PRs (30d)
- 25
Description
### Intro
Hi!
I am a undergrad student, I use MJX for my research on RL, mainly exploration methods. I encoutered an issue with my model and tried to reproduce it with a minimal example.
### My setup
- OS: macOS (darwin 24.3.0)
- Python: 3.12 (uv-managed env)
- MuJoCo: 3.3.7
- MJX: bundled with the above MuJoCo
- JAX: recent CPU build (default via MJX deps)
### What's happening? What did you expect?
What’s happening:
- In a minimal repro (single leg, only front-left lift tendon kept; other legs, walls, target, and equality removed):
- With MJX: mjx.forward succeeds, but the very first mjx.step immediately produces NaNs (qpos/qvel).
- With MuJoCo native: mj_forward and mj_step are stable (no NaNs).
- Enabling JAX float64 globally (config.update("jax_enable_x64", True) before importing mujoco.mjx) does not change the outcome: MJX still NaNs on the first step.
- Changing integrator and timestep (Euler/RK4; dt = 0.005/0.002/0.001) does not help: MJX still NaNs on the first step.
- Actuator/transmission ablations (general/motor, joint/tendon, general with dyntype=none + small gain + zero bias) all still NaN in MJX, while MuJoCo native stays stable.
- Replacing the freejoint root with a weld (fixing the torso to world) also still NaNs on the first step in MJX.
What did you expect:
- MJX step should be numerically stable for this model and initial state, matching MuJoCo native behavior.
- Even with rigid constraints or tendon/actuator combinations, MJX should not produce NaNs on the very first simulation step.
### Steps for reproduction
1. Load the model below.
2. Run the code below.
3. Should See:
forward NaN: qpos=False, qvel=False
step 1: NaN=True
MuJoCo forward NaN: qpos=False, qvel=False
MuJoCo step 1: NaN=False
MuJoCo step 2: NaN=False
MuJoCo step 3: NaN=False
MuJoCo step 4: NaN=False
MuJoCo step 5: NaN=False
### Minimal model for reproduction
minimal XML
```XML
```
### Code required for reproduction
```python
import jax
from jax import config
config.update("jax_enable_x64", True)
import mujoco
import mujoco.mjx as mjx
import jax.numpy as jp
import numpy as np
XML = "./quad_single_leg.xml"
XML = "/Users/dopamine/PJRFE/mjx-cov/repro_min/quad_single_leg.xml"
mj_model = mujoco.MjModel.from_xml_path(XML)
mjx_model = mjx.put_model(mj_model)
data = mjx.make_data(mjx_model)
qpos = jp.zeros(mjx_model.nq)
qpos = qpos.at[2].set(0.5)
qpos = qpos.at[3:7].set(jp.array([1.0, 0.0, 0.0, 0.0]))
qvel = jp.zeros(mjx_model.nv)
ctrl = jp.zeros(mjx_model.nu)
data = data.replace(qpos=qpos, qvel=qvel, ctrl=ctrl)
data = mjx.forward(mjx_model, data)
print(f"forward NaN: qpos={bool(jp.isnan(data.qpos).any())}, qvel={bool(jp.isnan(data.qvel).any())}")
@jax.jit
def step_jit(m, d):
d = mjx.forward(m, d)
d = mjx.step(m, d)
return d
nan=False
for i in range(5):
data = step_jit(mjx_model, data)
if bool(jp.isnan(data.qpos).any() or jp.isnan(data.qvel).any()):
print(f"step {i+1}: NaN=True")
nan=True
break
else:
print(f"step {i+1}: NaN=False")
md = mujoco.MjData(mj_model)
md.qpos[:] = 0.0
md.qpos[2] = 0.5
md.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
md.qvel[:] = 0.0
md.ctrl[:] = 0.0
mujoco.mj_forward(mj_model, md)
print(f"MuJoCo forward NaN: qpos={np.isnan(md.qpos).any()}, qvel={np.isnan(md.qvel).any()}")
for i in range(5):
mujoco.mj_step(mj_model, md)
has_nan = np.isnan(md.qpos).any() or np.isnan(md.qvel).any()
print(f"MuJoCo step {i+1}: NaN={has_nan}")
```
### Confirmations
- [x] I searched the [latest documentation](https://mujoco.readthedocs.io/en/latest/overview.html) thoroughly before posting.
- [x] I searched previous [Issues](https://github.com/google-deepmind/mujoco/issues) and [Discussions](https://github.com/google-deepmind/mujoco/discussions), I am certain this has not been raised before.
Contributor guide
Assessment
This issue has not been assessed yet.