add touch sensor

mujoco_warp/_src/io.py

@@ -627,7 +627,15 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     ndim=1,
   )
   m.sensor_acc_adr = wp.array(
-    np.nonzero(mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC)[0],
+    np.nonzero(
+      (mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC)
+      & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TOUCH)
+    )[0],
+    dtype=wp.int32,
+    ndim=1,
+  )
+  m.sensor_touch_adr = wp.array(
+    np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_TOUCH)[0],
     dtype=wp.int32,
     ndim=1,
   )

mujoco_warp/_src/sensor.py

@@ -21,6 +21,7 @@
 from . import math
 from . import smooth
 from .types import MJ_MINVAL
+from .types import ConeType
 from .types import Data
 from .types import DataType
 from .types import DisableBit
@@ -739,6 +740,60 @@ def _frameangacc(m: Model, d: Data, worldid: int, objid: int, objtype: int) -> w
 def sensor_acc(m: Model, d: Data):
   """Compute acceleration-dependent sensor values."""
 
+  @kernel
+  def _sensor_touch_zero(m: Model, d: Data):
+    worldid, sensortouchadrid = wp.tid()
+    sensorid = m.sensor_touch_adr[sensortouchadrid]
+    adr = m.sensor_adr[sensorid]
+    d.sensordata[worldid, adr] = 0.0
+
+  @kernel
+  def _sensor_touch(m: Model, d: Data):
+    conid, sensortouchadrid = wp.tid()
+
+    if conid > d.ncon[0]:
+      return
+
+    sensorid = m.sensor_touch_adr[sensortouchadrid]
+
+    objid = m.sensor_objid[sensorid]
+    bodyid = m.site_bodyid[objid]
+
+    # find contact in sensor zone, add normal force
+
+    # contacting bodies
+    geom = d.contact.geom[conid]
+    conbody = wp.vec2i(m.geom_bodyid[geom[0]], m.geom_bodyid[geom[1]])
+
+    # select contacts involving sensorized body
+    efc_address0 = d.contact.efc_address[conid, 0]
+    if efc_address0 >= 0 and (bodyid == conbody[0] or bodyid == conbody[1]):
+      # get contact normal force
+      normalforce = d.efc.force[efc_address0]
+
+      if m.opt.cone == int(ConeType.PYRAMIDAL.value):
+        dim = d.contact.dim[conid]
+        for i in range(1, 2 * (dim - 1)):
+          normalforce += d.efc.force[d.contact.efc_address[conid, i]]
+
+      if normalforce <= 0.0:
+        return
+
+      # convert contact normal force to global frame, normalize
+      frame = d.contact.frame[conid]
+      conray = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2]) * normalforce
+      conray, _ = math.normalize_with_norm(conray)
+
+      # flip ray direction if sensor is on body2
+      if bodyid == conbody[1]:
+        conray = -conray
+
+      # add if ray-zone intersection (always true when contact.pos inside zone)
+      worldid = d.contact.worldid[conid]
+      if True:  # TODO(team): ray.ray_geom(d.site_xpos[worldid, objid], d.site_xmat[worldid, objid], m.site_size[objid], d.contact.pos[conid], conray, m.site_type[objid]) >= 0.0
+        adr = m.sensor_adr[sensorid]
+        wp.atomic_add(d.sensordata[worldid], adr, normalforce)
+
   @kernel
   def _sensor_acc(m: Model, d: Data):
     worldid, accid = wp.tid()
@@ -788,10 +843,17 @@ def _sensor_acc(m: Model, d: Data):
       d.sensordata[worldid, adr + 1] = frame_angacc[1]
       d.sensordata[worldid, adr + 2] = frame_angacc[2]
 
-  if (m.sensor_acc_adr.size == 0) or (m.opt.disableflags & DisableBit.SENSOR):
+  if m.opt.disableflags & DisableBit.SENSOR:
     return
 
-  if m.sensor_rne_postconstraint:
-    smooth.rne_postconstraint(m, d)
+  if m.sensor_touch_adr.size > 0:
+    wp.launch(
+      _sensor_touch_zero, dim=(d.nworld, m.sensor_touch_adr.size), inputs=[m, d]
+    )
+    wp.launch(_sensor_touch, dim=(d.nconmax, m.sensor_touch_adr.size), inputs=[m, d])
+
+  if m.sensor_acc_adr.size > 0:
+    if m.sensor_rne_postconstraint:
+      smooth.rne_postconstraint(m, d)
 
-  wp.launch(_sensor_acc, dim=(d.nworld, m.sensor_acc_adr.size), inputs=[m, d])
+    wp.launch(_sensor_acc, dim=(d.nworld, m.sensor_acc_adr.size), inputs=[m, d])

mujoco_warp/_src/sensor_test.py

@@ -200,6 +200,50 @@ def test_sensor(self):
 
     _assert_eq(d.sensordata.numpy()[0], mjd.sensordata, "sensordata")
 
+  def test_touch_sensor(self):
+    """Test touch sensor."""
+    for keyframe in range(2):
+      _, mjd, m, d = test_util.fixture(
+        xml="""
+        <mujoco>
+          <worldbody>
+            <geom type="plane" size="10 10 .001"/>
+            <body pos="0 0 .25">
+              <geom type="sphere" size="0.1" pos=".1 0 0"/>
+              <geom type="sphere" size="0.1" pos="-.1 0 0"/>
+              <geom type="sphere" size="0.1" pos="0 0 .11"/>
+              <geom type="sphere" size="0.1" pos="0 0 10"/>
+              <site name="site_sphere" type="sphere" size=".2"/>
+              <site name="site_capsule" type="capsule" size=".2 .2"/>
+              <site name="site_ellipsoid" type="ellipsoid" size=".2 .2 .2"/>
+              <site name="site_cylinder" type="cylinder" size=".2 .2"/>
+              <site name="site_box" type="box" size=".2 .2 .2"/>
+              <freejoint/>
+            </body>
+          </worldbody>
+          <sensor>
+            <touch site="site_sphere"/>
+            <touch site="site_capsule"/>
+            <touch site="site_ellipsoid"/>
+            <touch site="site_cylinder"/>
+            <touch site="site_box"/>
+          </sensor>
+          <keyframe>
+            <key qpos="0 0 10 1 0 0 0"/>
+            <key qpos="0 0 .05 1 0 0 0"/>
+            <key qpos="0 0 0 1 0 0 0"/>
+          </keyframe>
+        </mujoco>
+      """,
+        keyframe=keyframe,
+      )
+
+      d.sensordata.zero_()
+
+      mjwarp.sensor_acc(m, d)
+
+      _assert_eq(d.sensordata.numpy()[0], mjd.sensordata, "sensordata")
+
   def test_tendon_sensor(self):
     """Test tendon sensors."""
     _, mjd, m, d = test_util.fixture("tendon/fixed.xml", keyframe=0, sparse=False)

mujoco_warp/_src/types.py

@@ -261,6 +261,7 @@ class SensorType(enum.IntEnum):
     FRAMEANGVEL: 3D angular velocity
     SUBTREELINVEL: subtree linear velocity
     SUBTREEANGMOM: subtree angular momentum
+    TOUCH: scalar contact normal forces summed over sensor zone
     ACCELEROMETER: accelerometer
     FORCE: force
     TORQUE: torque
@@ -292,6 +293,7 @@ class SensorType(enum.IntEnum):
   FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL
   SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL
   SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM
+  TOUCH = mujoco.mjtSensor.mjSENS_TOUCH
   ACCELEROMETER = mujoco.mjtSensor.mjSENS_ACCELEROMETER
   FORCE = mujoco.mjtSensor.mjSENS_FORCE
   TORQUE = mujoco.mjtSensor.mjSENS_TORQUE
@@ -761,6 +763,8 @@ class Model:
     sensor_pos_adr: addresses for position sensors           (<=nsensor,)
     sensor_vel_adr: addresses for velocity sensors           (<=nsensor,)
     sensor_acc_adr: addresses for acceleration sensors       (<=nsensor,)
+                    (excluding touch sensors)
+    sensor_touch_adr: addresses for touch sensors            (<=nsensor,)
     sensor_subtree_vel: evaluate subtree_vel
     sensor_rne_postconstraint: evaluate rne_postconstraint
   """
@@ -983,6 +987,7 @@ class Model:
   sensor_pos_adr: wp.array(dtype=wp.int32, ndim=1)  # warp only
   sensor_vel_adr: wp.array(dtype=wp.int32, ndim=1)  # warp only
   sensor_acc_adr: wp.array(dtype=wp.int32, ndim=1)  # warp only
+  sensor_touch_adr: wp.array(dtype=wp.int32, ndim=1)  # warp only
   sensor_subtree_vel: bool  # warp only
   sensor_rne_postconstraint: bool  # warp only