From 8be966cdf9073813ec8b494062f4d97848432057 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Mon, 29 Jan 2024 12:08:35 -0800 Subject: [PATCH] Avoid NaN in gradient of where. Fixes #1349. PiperOrigin-RevId: 602461299 Change-Id: I7c976d3b7262027bdc63967dcc00674fb2802e70 --- mjx/mujoco/mjx/_src/smooth.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 9d1d0c37..3d7632d7 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -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