From 0ba31e2c1fea3e9b5492977f25feff2be34df113 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 14 Aug 2025 05:48:22 -0700 Subject: [PATCH] MJX safe division. Fixes #2776, #2568, #2657. PiperOrigin-RevId: 794997676 Change-Id: I3100af7c1987a8435b29cc68c4e5c624dc573986 --- mjx/mujoco/mjx/_src/collision_primitive.py | 2 +- mjx/mujoco/mjx/_src/math.py | 8 ++++++ mjx/mujoco/mjx/_src/ray.py | 14 +++++------ mjx/mujoco/mjx/_src/smooth.py | 11 +++++--- mjx/mujoco/mjx/_src/support.py | 29 ++++++++++++---------- 5 files changed, 40 insertions(+), 24 deletions(-) diff --git a/mjx/mujoco/mjx/_src/collision_primitive.py b/mjx/mujoco/mjx/_src/collision_primitive.py index eebd3eb3..0677dd37 100644 --- a/mjx/mujoco/mjx/_src/collision_primitive.py +++ b/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 diff --git a/mjx/mujoco/mjx/_src/math.py b/mjx/mujoco/mjx/_src/math.py index bd3b6019..e52d41ab 100644 --- a/mjx/mujoco/mjx/_src/math.py +++ b/mjx/mujoco/mjx/_src/math.py @@ -18,6 +18,14 @@ from typing import Optional, Tuple, Union 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: diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py index 88db31be..2e2055cf 100644 --- a/mjx/mujoco/mjx/_src/ray.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 1303f357..dd6aa89e 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/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 tendon(m: Model, d: Data) -> Data: 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 tendon_dot(m: Model, d: Data) -> jax.Array: 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) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index ba62d9e3..04b0a0d4 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -470,7 +470,12 @@ class BindData(object): 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 wrap_inside( 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 muscle_dynamics_timescale( # 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)