chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Avoid NaN in gradient of where. Fixes #1349.

Browse Source 
PiperOrigin-RevId: 602461299
Change-Id: I7c976d3b7262027bdc63967dcc00674fb2802e70
This commit is contained in:
Baruch Tabanpour
2024-01-29 12:08:35 -08:00
committed by Copybara-Service
parent 508669a9f2
commit 8be966cdf9
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/smooth.py
+3 -1
View File
@@ -137,7 +137,9 @@ def com_pos(m: Model, d: Data) -> Data:
m, subtree_sum, 'bb', 'bb', d.xipos, m.body_mass, reverse=True m, subtree_sum, 'bb', 'bb', d.xipos, m.body_mass, reverse=True
) )
cond = jp.tile(mass < mujoco.mjMINVAL, (3, 1)).T 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) d = d.replace(subtree_com=subtree_com)
# map inertias to frame centered at subtree_com # map inertias to frame centered at subtree_com