Avoid NaN in gradient of where. Fixes #1349.

PiperOrigin-RevId: 602461299 Change-Id: I7c976d3b7262027bdc63967dcc00674fb2802e70
2024-01-29 12:08:35 -08:00
parent 508669a9f2
commit 8be966cdf9
1 changed files with 3 additions and 1 deletions
@@ -137,7 +137,9 @@ def com_pos(m: Model, d: Data) -> Data:
      m, subtree_sum, 'bb', 'bb', d.xipos, m.body_mass, reverse=True
  )
  cond = jp.tile(mass < mujoco.mjMINVAL, (3, 1)).T
-  subtree_com = jp.where(cond, d.xipos, jax.vmap(jp.divide)(pos, mass))
+  # take maximum to avoid NaN in gradient of jp.where
+  subtree_com = jax.vmap(jp.divide)(pos, jp.maximum(mass, mujoco.mjMINVAL))
+  subtree_com = jp.where(cond, d.xipos, subtree_com)
  d = d.replace(subtree_com=subtree_com)

  # map inertias to frame centered at subtree_com