contact sensor maxforce reduction (#602)

* contact sensor maxforce reduction

* update test_contact_sensor

* update test_contact_sensor

* update test_contact_sensor to test fewer contact sensor instances

Author: thowell

mujoco_warp/_src/io.py

@@ -129,8 +129,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
       raise NotImplementedError("Contact sensor: only geom1-geom2 matching is implemented.")
 
     # reduction
-    if (mjm.sensor_intprm[is_contact_sensor, 1] != 1).any():
-      raise NotImplementedError(f"Contact sensor: only mindist reduction is implemented.")
+    if (~((mjm.sensor_intprm[is_contact_sensor, 1] == 1) | (mjm.sensor_intprm[is_contact_sensor, 1] == 2))).any():
+      raise NotImplementedError(f"Contact sensor: only mindist and maxforce reduction are implemented.")
 
   # TODO(team): remove after _update_gradient for Newton uses tile operations for islands
   nv_max = 60

mujoco_warp/_src/io_test.py

@@ -366,7 +366,6 @@ def test_make_put_data_dims_match(self, xml):
     '<contact geom1="plane"/>',
     '<contact geom2="plane"/>',
     '<contact site="site"/>',
-    '<contact reduce="maxforce"/>',
     '<contact reduce="netforce"/>',
     '<contact geom1="plane" geom2="sphere"/>',
   )

mujoco_warp/_src/sensor.py

@@ -1986,14 +1986,22 @@ def _sensor_tactile(
 @wp.kernel
 def _contact_match(
   # Model:
+  opt_cone: int,
   sensor_objid: wp.array(dtype=int),
   sensor_refid: wp.array(dtype=int),
+  sensor_intprm: wp.array2d(dtype=int),
   sensor_contact_adr: wp.array(dtype=int),
   # Data in:
+  njmax_in: int,
   ncon_in: wp.array(dtype=int),
   contact_dist_in: wp.array(dtype=float),
+  contact_frame_in: wp.array(dtype=wp.mat33),
+  contact_friction_in: wp.array(dtype=vec5),
+  contact_dim_in: wp.array(dtype=int),
   contact_geom_in: wp.array(dtype=wp.vec2i),
+  contact_efc_address_in: wp.array2d(dtype=int),
   contact_worldid_in: wp.array(dtype=int),
+  efc_force_in: wp.array2d(dtype=float),
   # Data out:
   sensor_contact_nmatch_out: wp.array2d(dtype=int),
   sensor_contact_matchid_out: wp.array3d(dtype=int),
@@ -2009,6 +2017,7 @@ def _contact_match(
   # sensor information
   objid = sensor_objid[sensorid]
   refid = sensor_refid[sensorid]
+  reduce = sensor_intprm[sensorid, 1]
 
   # contact information
   geom = contact_geom_in[contactid]
@@ -2022,8 +2031,27 @@ def _contact_match(
     contactmatchid = wp.atomic_add(sensor_contact_nmatch_out[worldid], contactsensorid, 1)
     sensor_contact_matchid_out[worldid, contactsensorid, contactmatchid] = contactid
 
-    # TODO(thowell): alternative criteria
-    sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = contact_dist_in[contactid]
+    if reduce == 1:  # mindist
+      sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = contact_dist_in[contactid]
+    elif reduce == 2:  # maxforce
+      contact_force = support.contact_force_fn(
+        opt_cone,
+        njmax_in,
+        ncon_in,
+        contact_frame_in,
+        contact_friction_in,
+        contact_dim_in,
+        contact_efc_address_in,
+        efc_force_in,
+        worldid,
+        contactid,
+        False,
+      )
+      force_magnitude = (
+        contact_force[0] * contact_force[0] + contact_force[1] * contact_force[1] + contact_force[2] * contact_force[2]
+      )
+      sensor_contact_criteria_out[worldid, contactsensorid, contactmatchid] = -force_magnitude
+    # TODO(thowell): netforce
 
     # contact direction
     if geom1geom0:
@@ -2155,20 +2183,28 @@ def sensor_acc(m: Model, d: Data):
   if m.sensor_contact_adr.size:
     # match criteria
     d.sensor_contact_nmatch.zero_()
-    d.sensor_contact_matchid.zero_()
-    d.sensor_contact_criteria.zero_()
+    d.sensor_contact_matchid.fill_(-1)
+    d.sensor_contact_criteria.fill_(1.0e32)
 
     wp.launch(
       _contact_match,
       dim=(m.sensor_contact_adr.size, d.nconmax),
       inputs=[
+        m.opt.cone,
         m.sensor_objid,
         m.sensor_refid,
+        m.sensor_intprm,
         m.sensor_contact_adr,
+        d.njmax,
         d.ncon,
         d.contact.dist,
+        d.contact.frame,
+        d.contact.friction,
+        d.contact.dim,
         d.contact.geom,
+        d.contact.efc_address,
         d.contact.worldid,
+        d.efc.force,
       ],
       outputs=[
         d.sensor_contact_nmatch,

mujoco_warp/_src/sensor_test.py

@@ -411,30 +411,26 @@ def test_energy(self, xml):
   @parameterized.parameters(
     'type="sphere" size=".1"',
     'type="capsule" size=".1 .1" euler="0 89 89"',
-    'type="box" size=".1 .1 .1" euler=".02 .05 .1"',
+    'type="box" size=".1 .11 .12" euler=".02 .05 .1"',
   )
   def test_contact_sensor(self, geom):
     """Test contact sensor."""
     # create contact sensors
     contact_sensor = ""
 
     # data combinations
-    field = ["found", "force", "torque", "dist", "pos", "normal", "tangent"]
-    datas = itertools.chain.from_iterable([itertools.combinations(field, i) for i in range(len(field))])
-    datas = list(datas)
-
-    for num in [1, 2, 3, 4, 5]:
-      for geoms in [
-        'geom1="plane" geom2="geom"',
-        'geom1="geom" geom2="plane"',
-        'geom1="plane" geom2="sphere"',
-        'geom1="sphere" geom2="plane"',
-        'geom1="geom" geom2="sphere"',
-        'geom1="sphere" geom2="geom"',
-      ]:
-        for data in datas:
-          data = " ".join(data)
-          contact_sensor += f'<contact {geoms} num="{num}" reduce="mindist" data="{data}"/>'
+    datas = ["found", "force dist normal", "torque pos tangent", "found force torque dist pos normal tangent"]
+
+    for geoms in [
+      'geom1="plane" geom2="geom"',
+      'geom1="geom" geom2="plane"',
+      'geom1="sphere" geom2="plane"',
+      'geom1="geom" geom2="sphere"',
+    ]:
+      for num in [1, 3, 5]:
+        for reduce in ["mindist", "maxforce"]:
+          for data in datas:
+            contact_sensor += f'<contact {geoms} num="{num}" reduce="{reduce}" data="{data}"/>'
 
     _MJCF = f"""
     <mujoco>