踩坑/GitHub google-deepmind/mujoco·9 天前

MJX spatial tendon内部sidesite导致jax.grad返回NaN

MJX中spatial tendon的sidesite位于wrap geom内部时梯度为NaN，需对Newton迭代及反三角函数输入增加数值安全钳位与条件替换。

AI 整理

• 仿真与训练

本文由 AI 根据公开内容整理，可能有误，请以原文为准：[MJX] jax.grad returns NaN through a spatial tendon with an inside-sidesite geom wrap: wrap_inside's Newton iterate leaves arcsin's domain and the NaN survives the jp.where（GitHub google-deepmind/mujoco #3696）

状态：有绕过办法

问题

在MuJoCo MJX中，当spatial tendon的geom wrap包含位于包裹几何体内部的sidesite时，jax.grad计算ten_length梯度静默返回NaN（前向值正确）。根因是wrap_inside的Newton迭代使arcsin参数超出定义域，且support.wrap中圆柱校正的sqrt(0)导数无穷大，以及退化姿态下acos(±1)不可微，导致自动微分失败。

环境

MuJoCo/MJX 3.16.0, JAX 0.11.2, macOS arm64 CPU, float32精度

解决办法

1. 在support.wrap_inside的Newton更新后添加z_next = jp.clip(z_next, -1.0 + 1e-7, 1.0 - 1e-7)防止arcsin越界。
2. 将support.wrap中的jp.sqrt替换为math.norm以处理零向量处的无限导数问题。
3. 在调用jp.arccos(cosG)前，针对float32精度定义安全掩码：no_wrap2 = (cosG < -1 + mjMINVAL) | (cosG <= -1)与early_return0 = (cosG > 1 - mjMINVAL) | (cosG >= 1)。
4. 合并所有提前退出条件skip = no_wrap0 | no_wrap1 | no_wrap2 | early_return0。
5. 使用jp.where(skip, safe_val, original_val)对A、B（替换为0.5）及cosG（替…

讨论原作者：@adam-hartshorne 等，版权归原作者所有。

来源

本页事实对应的公开页面，以原文为准。

1. GitHub google-deepmind/mujoco #3696：[MJX] jax.grad returns NaN through a spatial tendon with an inside-sidesite geom wrap: wrap_inside's Newton iterate leaves arcsin's domain and the NaN survives the jp.where访问 2026-10-26