MJX safe division. Fixes #2776, #2568, #2657.

PiperOrigin-RevId: 794997676
Change-Id: I3100af7c1987a8435b29cc68c4e5c624dc573986

Author: thowell

mjx/mujoco/mjx/_src/collision_primitive.py

@@ -120,7 +120,7 @@ def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Collision:
       # disk parallel to plane: pick x-axis of cylinder, scale by radius
       cylinder.mat[:, 0] * cylinder.size[0],
       # general configuration: normalize vector, scale by radius
-      vec / len_ * cylinder.size[0],
+      math.safe_div(vec, len_) * cylinder.size[0],
   )
 
   # project vector on normal

mjx/mujoco/mjx/_src/math.py

@@ -18,6 +18,14 @@
 
 import jax
 from jax import numpy as jp
+import mujoco
+
+
+def safe_div(
+    num: Union[float, jax.Array], den: Union[float, jax.Array]
+) -> Union[float, jax.Array]:
+  """Safe division for case where denominator is zero."""
+  return num / (den + mujoco.mjMINVAL * (den == 0))
 
 
 def matmul_unroll(a: jax.Array, b: jax.Array) -> jax.Array:

mjx/mujoco/mjx/_src/ray.py

@@ -35,7 +35,7 @@ def _ray_quad(
   det = b * b - a * c
   det_2 = jp.sqrt(det)
 
-  x0, x1 = (-b - det_2) / a, (-b + det_2) / a
+  x0, x1 = math.safe_div(-b - det_2, a), math.safe_div(-b + det_2, a)
   x0 = jp.where((det < mujoco.mjMINVAL) | (x0 < 0), jp.inf, x0)
   x1 = jp.where((det < mujoco.mjMINVAL) | (x1 < 0), jp.inf, x1)
 
@@ -48,7 +48,7 @@ def _ray_plane(
     vec: jax.Array,
 ) -> jax.Array:
   """Returns the distance at which a ray intersects with a plane."""
-  x = -pnt[2] / vec[2]
+  x = -math.safe_div(pnt[2], vec[2])
 
   valid = vec[2] <= -mujoco.mjMINVAL  # z-vec pointing towards front face
   valid &= x >= 0
@@ -116,7 +116,7 @@ def _ray_ellipsoid(
   """Returns the distance at which a ray intersects with an ellipsoid."""
 
   # invert size^2
-  s = 1 / jp.square(size)
+  s = math.safe_div(1, jp.square(size))
 
   # (x*lvec+lpnt)' * diag(1/size^2) * (x*lvec+lpnt) = 1
   svec = s * vec
@@ -142,7 +142,7 @@ def _ray_box(
 
   # side +1, -1
   # solution of pnt[i] + x * vec[i] = side * size[i]
-  x = jp.concatenate([(size - pnt) / vec, (-size - pnt) / vec])
+  x = jp.concatenate([math.safe_div(size - pnt,  vec), -math.safe_div(size + pnt, vec)])
 
   # intersection with face
   p0 = pnt[iface[:, 0]] + x * vec[iface[:, 0]]
@@ -170,13 +170,13 @@ def _ray_triangle(
   b = -planar[2]
   det = A[0, 0] * A[1, 1] - A[1, 0] * A[0, 1]
 
-  t0 = (A[1, 1] * b[0] - A[1, 0] * b[1]) / det
-  t1 = (-A[0, 1] * b[0] + A[0, 0] * b[1]) / det
+  t0 = math.safe_div(A[1, 1] * b[0] - A[1, 0] * b[1], det)
+  t1 = math.safe_div(-A[0, 1] * b[0] + A[0, 0] * b[1], det)
   valid = (t0 >= 0) & (t1 >= 0) & (t0 + t1 <= 1)
 
   # intersect ray with plane of triangle
   nrm = jp.cross(vert[0] - vert[2], vert[1] - vert[2])
-  dist = jp.dot(vert[2] - pnt, nrm) / jp.dot(vec, nrm)
+  dist = math.safe_div(jp.dot(vert[2] - pnt, nrm), jp.dot(vec, nrm))
   valid &= dist >= 0
   dist = jp.where(valid, dist, jp.inf)
 

mjx/mujoco/mjx/_src/smooth.py

@@ -44,6 +44,7 @@ def kinematics(m: Model, d: Data) -> Data:
     from mujoco.mjx.warp import smooth as mjxw_smooth  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
     return mjxw_smooth.kinematics(m, d)
 
+
   def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat):
     # calculate joint anchors, axes, body pos and quat in global frame
     # also normalize qpos while we're at it
@@ -910,7 +911,9 @@ def _length_moment(pnt0, pnt1, body0, body1):
     dif = pnt1 - pnt0
     length = math.norm(dif)
     vec = jp.where(
-        length < mujoco.mjMINVAL, jp.array([1.0, 0.0, 0.0]), dif / length
+        length < mujoco.mjMINVAL,
+        jp.array([1.0, 0.0, 0.0]),
+        math.safe_div(dif, length),
     )
 
     jacp1, _ = support.jac(m, d, pnt0, body0)
@@ -1387,14 +1390,16 @@ def _momentdot(wpnt0, wpnt1, wvel0, wvel1, body0, body1):
     dpnt = wpnt1 - wpnt0
     norm = math.norm(dpnt)
     dpnt = jp.where(
-        norm < mujoco.mjMINVAL, jp.array([1.0, 0.0, 0.0]), dpnt / norm
+        norm < mujoco.mjMINVAL,
+        jp.array([1.0, 0.0, 0.0]),
+        math.safe_div(dpnt, norm),
     )
 
     # dvel = d / dt(dpnt)
     dvel = wvel1 - wvel0
     dot = jp.dot(dpnt, dvel)
     dvel += dpnt * -dot
-    dvel = jp.where(norm > mujoco.mjMINVAL, dvel / norm, 0.0)
+    dvel = jp.where(norm > mujoco.mjMINVAL, math.safe_div(dvel, norm), 0.0)
 
     # get endpoint JacobianDots, subtract
     jacp1, _ = support.jac_dot(m, d, wpnt0, body0)

mjx/mujoco/mjx/_src/support.py

@@ -470,7 +470,12 @@ def __getname(self, name: str):
         return name
       else:
         raise AttributeError('ctrl is not available for this type')
-    if name == 'qpos' or name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
+    if (
+        name == 'qpos'
+        or name == 'qvel'
+        or name == 'qacc'
+        or name.startswith('qfrc_')
+    ):
       if self.prefix == 'jnt_':
         return name
       else:
@@ -672,12 +677,12 @@ def _is_intersect(
   det = (p4[1] - p3[1]) * (p2[0] - p1[0]) - (p4[0] - p3[0]) * (p2[1] - p1[1])
 
   # compute intersection point on each line
-  a = (
-      (p4[0] - p3[0]) * (p1[1] - p3[1]) - (p4[1] - p3[1]) * (p1[0] - p3[0])
-  ) / det
-  b = (
-      (p2[0] - p1[0]) * (p1[1] - p3[1]) - (p2[1] - p1[1]) * (p1[0] - p3[0])
-  ) / det
+  a = math.safe_div(
+      (p4[0] - p3[0]) * (p1[1] - p3[1]) - (p4[1] - p3[1]) * (p1[0] - p3[0]), det
+  )
+  b = math.safe_div(
+      (p2[0] - p1[0]) * (p1[1] - p3[1]) - (p2[1] - p1[1]) * (p1[0] - p3[0]), det
+  )
 
   return jp.where(
       jp.abs(det) < mujoco.mjMINVAL,
@@ -856,9 +861,7 @@ def _newton(carry, _):
     status0 = df > -mjMINVAL
 
     # new point
-    z_next = z - (1 - converged) * f / jp.where(
-        jp.abs(df) < mjMINVAL, mjMINVAL, df
-    )
+    z_next = z - (1 - converged) * math.safe_div(f, df)
 
     # make sure we are moving to the left; SHOULD NOT OCCUR
     status1 = z_next > z
@@ -987,8 +990,8 @@ def wrap(
   l1 = jp.sqrt(
       (p1[0] - res[3]) * (p1[0] - res[3]) + (p1[1] - res[4]) * (p1[1] - res[4])
   )
-  r2 = p0[2] + (p1[2] - p0[2]) * l0 / (l0 + wlen + l1)
-  r5 = p0[2] + (p1[2] - p0[2]) * (l0 + wlen) / (l0 + wlen + l1)
+  r2 = p0[2] + (p1[2] - p0[2]) * math.safe_div(l0, l0 + wlen + l1)
+  r5 = p0[2] + (p1[2] - p0[2]) * math.safe_div(l0 + wlen, l0 + wlen + l1)
   height = jp.abs(r5 - r2)
 
   wlen = jp.where(is_sphere, wlen, jp.sqrt(wlen * wlen + height * height))
@@ -1130,7 +1133,7 @@ def _sigmoid(x):
   # smooth switching
   # scale by width, center around 0.5 midpoint, rescale to bounds
   tau_smooth = tau_deact + (tau_act - tau_deact) * _sigmoid(
-      dctrl / smoothing_width + 0.5
+      math.safe_div(dctrl, smoothing_width) + 0.5
   )
 
   return jp.where(smoothing_width < mujoco.mjMINVAL, tau_hard, tau_smooth)