Replace margin/gap max with sum.

PiperOrigin-RevId: 868842521
Change-Id: Ia4a4ca5c6821c5a7b3ba7b0af7f71d981c2866c1

mjx/mujoco/mjx/_src/collision_driver.py

@@ -283,8 +283,8 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
       ))
     if geom1.size > 0 and geom2.size > 0:
       # other contacts get their params from geom fields
-      margin = jp.maximum(m.geom_margin[geom1], m.geom_margin[geom2])
-      gap = jp.maximum(m.geom_gap[geom1], m.geom_gap[geom2])
+      margin = m.geom_margin[geom1] + m.geom_margin[geom2]
+      gap = m.geom_gap[geom1] + m.geom_gap[geom2]
       solmix1, solmix2 = m.geom_solmix[geom1], m.geom_solmix[geom2]
       mix = solmix1 / (solmix1 + solmix2)
       mix = jp.where((solmix1 < eps) & (solmix2 < eps), 0.5, mix)

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py

@@ -91,7 +91,7 @@ def _hfield_filter(
   r2 = geom_rbound[rbound_id, g2]
 
   # TODO(team): margin?
-  margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
+  margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
 
   # box-sphere test: horizontal plane
   for i in range(2):

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py

@@ -101,18 +101,18 @@ def _plane_filter(
   if size1 == 0.0:
     # geom1 is a plane
     dist = wp.dot(xpos2 - xpos1, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2]))
-    return dist <= size2 + wp.max(margin1, margin2)
+    return dist <= size2 + margin1 + margin2
   elif size2 == 0.0:
     # geom2 is a plane
     dist = wp.dot(xpos1 - xpos2, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2]))
-    return dist <= size1 + wp.max(margin1, margin2)
+    return dist <= size1 + margin1 + margin2
 
   return True
 
 
 @wp.func
 def _sphere_filter(size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3) -> bool:
-  bound = size1 + size2 + wp.max(margin1, margin2)
+  bound = size1 + size2 + margin1 + margin2
   dif = xpos2 - xpos1
   dist_sq = wp.dot(dif, dif)
   return dist_sq <= bound * bound
@@ -141,7 +141,7 @@ def _aabb_filter(
   center1 = xmat1 @ center1 + xpos1
   center2 = xmat2 @ center2 + xpos2
 
-  margin = wp.max(margin1, margin2)
+  margin = margin1 + margin2
 
   max_x1 = -MJ_MAXVAL
   max_y1 = -MJ_MAXVAL
@@ -236,7 +236,7 @@ def _obb_filter(
   xmat2: wp.mat33,
 ) -> bool:
   """Oriented bounding boxes collision (see Gottschalk et al.), see mj_collideOBB."""
-  margin = wp.max(margin1, margin2)
+  margin = margin1 + margin2
 
   xcenter = mat23()
   normal = mat63()

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py

@@ -552,8 +552,8 @@ def contact_params(
     solreffriction = wp.vec2(0.0, 0.0)
     solimp = mix * geom_solimp[solimp_id, g1] + (1.0 - mix) * geom_solimp[solimp_id, g2]
     # geom priority is ignored
-    margin = wp.max(geom_margin[margin_id, g1], geom_margin[margin_id, g2])
-    gap = wp.max(geom_gap[gap_id, g1], geom_gap[gap_id, g2])
+    margin = geom_margin[margin_id, g1] + geom_margin[margin_id, g2]
+    gap = geom_gap[gap_id, g1] + geom_gap[gap_id, g2]
 
   friction = vec5(
     wp.max(MJ_MINMU, friction[0]),