add touch sensor

mujoco_warp/_src/io.py

@@ -405,7 +405,9 @@ def create_nmodel_batched_array(mjm_array, dtype):
     geom_margin=create_nmodel_batched_array(mjm.geom_margin, dtype=float),
     geom_gap=create_nmodel_batched_array(mjm.geom_gap, dtype=float),
     geom_rgba=create_nmodel_batched_array(mjm.geom_rgba, dtype=wp.vec4),
+    site_type=wp.array(mjm.site_type, dtype=int),
     site_bodyid=wp.array(mjm.site_bodyid, dtype=int),
+    site_size=wp.array(mjm.site_size, dtype=wp.vec3),
     site_pos=create_nmodel_batched_array(mjm.site_pos, dtype=wp.vec3),
     site_quat=create_nmodel_batched_array(mjm.site_quat, dtype=wp.quat),
     cam_mode=wp.array(mjm.cam_mode, dtype=int),
@@ -534,7 +536,11 @@ def create_nmodel_batched_array(mjm_array, dtype):
       dtype=int,
     ),
     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=int,
+    ),
+    sensor_touch_adr=wp.array(
+      np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_TOUCH)[0],
       dtype=int,
     ),
     sensor_subtree_vel=np.isin(

mujoco_warp/_src/ray.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from typing import Tuple
+
 import warp as wp
 
 from .types import MJ_MINVAL
@@ -361,7 +363,7 @@ def _ray_map(
   mat: wp.mat33,
   pnt: wp.vec3,
   vec: wp.vec3,
-) -> any:
+) -> Tuple[wp.vec3, wp.vec3]:
   """Maps ray to local geom frame coordinates.
 
   Args:
@@ -384,7 +386,7 @@ def _ray_map(
 
 
 @wp.func
-def _ray_geom(
+def ray_geom(
   # In:
   pos: wp.vec3,  # Position of geom frame
   mat: wp.mat33,  # Orientation of geom frame

mujoco_warp/_src/sensor.py

@@ -18,8 +18,10 @@
 import warp as wp
 
 from . import math
+from . import ray
 from . import smooth
 from .types import MJ_MINVAL
+from .types import ConeType
 from .types import Data
 from .types import DataType
 from .types import DisableBit
@@ -1273,44 +1275,179 @@ def _sensor_acc(
     _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
 
 
+@wp.kernel
+def _sensor_touch_zero(
+  # Model:
+  sensor_adr: wp.array(dtype=int),
+  sensor_touch_adr: wp.array(dtype=int),
+  # Data out:
+  sensordata_out: wp.array2d(dtype=float),
+):
+  worldid, sensortouchadrid = wp.tid()
+  sensorid = sensor_touch_adr[sensortouchadrid]
+  adr = sensor_adr[sensorid]
+  sensordata_out[worldid, adr] = 0.0
+
+
+@wp.kernel
+def _sensor_touch(
+  # Model:
+  opt_cone: int,
+  geom_bodyid: wp.array(dtype=int),
+  site_type: wp.array(dtype=int),
+  site_bodyid: wp.array(dtype=int),
+  site_size: wp.array(dtype=wp.vec3),
+  sensor_objid: wp.array(dtype=int),
+  sensor_adr: wp.array(dtype=int),
+  sensor_touch_adr: wp.array(dtype=int),
+  # Data in:
+  ncon_in: wp.array(dtype=int),
+  site_xpos_in: wp.array2d(dtype=wp.vec3),
+  site_xmat_in: wp.array2d(dtype=wp.mat33),
+  contact_pos_in: wp.array(dtype=wp.vec3),
+  contact_frame_in: wp.array(dtype=wp.mat33),
+  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.array(dtype=float),
+  # Data out:
+  sensordata_out: wp.array2d(dtype=float),
+):
+  conid, sensortouchadrid = wp.tid()
+
+  if conid > ncon_in[0]:
+    return
+
+  sensorid = sensor_touch_adr[sensortouchadrid]
+
+  objid = sensor_objid[sensorid]
+  bodyid = site_bodyid[objid]
+
+  # find contact in sensor zone, add normal force
+
+  # contacting bodies
+  geom = contact_geom_in[conid]
+  conbody = wp.vec2i(geom_bodyid[geom[0]], geom_bodyid[geom[1]])
+
+  # select contacts involving sensorized body
+  efc_address0 = contact_efc_address_in[conid, 0]
+  if efc_address0 >= 0 and (bodyid == conbody[0] or bodyid == conbody[1]):
+    # get contact normal force
+    normalforce = efc_force_in[efc_address0]
+
+    if opt_cone == int(ConeType.PYRAMIDAL.value):
+      dim = contact_dim_in[conid]
+      for i in range(1, 2 * (dim - 1)):
+        normalforce += efc_force_in[contact_efc_address_in[conid, i]]
+
+    if normalforce <= 0.0:
+      return
+
+    # convert contact normal force to global frame, normalize
+    frame = contact_frame_in[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 = contact_worldid_in[conid]
+    if (
+      ray.ray_geom(
+        site_xpos_in[worldid, objid],
+        site_xmat_in[worldid, objid],
+        site_size[objid],
+        contact_pos_in[conid],
+        conray,
+        site_type[objid],
+      )
+      >= 0.0
+    ):
+      adr = sensor_adr[sensorid]
+      wp.atomic_add(sensordata_out[worldid], adr, normalforce)
+
+
 @event_scope
 def sensor_acc(m: Model, d: Data):
   """Compute acceleration-dependent sensor values."""
-
-  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.sensor_adr,
+        m.sensor_touch_adr,
+      ],
+      outputs=[
+        d.sensordata,
+      ],
+    )
+    wp.launch(
+      _sensor_touch,
+      dim=(d.nconmax, m.sensor_touch_adr.size),
+      inputs=[
+        m.opt.cone,
+        m.geom_bodyid,
+        m.site_type,
+        m.site_bodyid,
+        m.site_size,
+        m.sensor_objid,
+        m.sensor_adr,
+        m.sensor_touch_adr,
+        d.ncon,
+        d.site_xpos,
+        d.site_xmat,
+        d.contact.pos,
+        d.contact.frame,
+        d.contact.dim,
+        d.contact.geom,
+        d.contact.efc_address,
+        d.contact.worldid,
+        d.efc.force,
+      ],
+      outputs=[
+        d.sensordata,
+      ],
+    )
 
-  wp.launch(
-    _sensor_acc,
-    dim=(d.nworld, m.sensor_acc_adr.size),
-    inputs=[
-      m.body_rootid,
-      m.jnt_dofadr,
-      m.geom_bodyid,
-      m.site_bodyid,
-      m.cam_bodyid,
-      m.sensor_type,
-      m.sensor_datatype,
-      m.sensor_objtype,
-      m.sensor_objid,
-      m.sensor_adr,
-      m.sensor_cutoff,
-      m.sensor_acc_adr,
-      d.xpos,
-      d.xipos,
-      d.geom_xpos,
-      d.site_xpos,
-      d.site_xmat,
-      d.cam_xpos,
-      d.subtree_com,
-      d.cvel,
-      d.actuator_force,
-      d.qfrc_actuator,
-      d.cacc,
-      d.cfrc_int,
-    ],
-    outputs=[d.sensordata],
-  )
+  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.body_rootid,
+        m.jnt_dofadr,
+        m.geom_bodyid,
+        m.site_bodyid,
+        m.cam_bodyid,
+        m.sensor_type,
+        m.sensor_datatype,
+        m.sensor_objtype,
+        m.sensor_objid,
+        m.sensor_adr,
+        m.sensor_cutoff,
+        m.sensor_acc_adr,
+        d.xpos,
+        d.xipos,
+        d.geom_xpos,
+        d.site_xpos,
+        d.site_xmat,
+        d.cam_xpos,
+        d.subtree_com,
+        d.cvel,
+        d.actuator_force,
+        d.qfrc_actuator,
+        d.cacc,
+        d.cfrc_int,
+      ],
+      outputs=[d.sensordata],
+    )

mujoco_warp/_src/sensor_test.py

@@ -233,6 +233,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

@@ -262,6 +262,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
@@ -293,6 +294,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
@@ -687,6 +689,7 @@ class Model:
     geom_margin: detect contact if dist<margin               (nworld, ngeom,)
     geom_gap: include in solver if dist<margin-gap           (nworld, ngeom,)
     geom_rgba: rgba when material is omitted                 (nworld, ngeom, 4)
+    site_type: geom type for rendering (mjtGeom)             (nsite,)
     site_bodyid: id of site's body                           (nsite,)
     site_pos: local position offset rel. to body             (nworld, nsite, 3)
     site_quat: local orientation offset rel. to body         (nworld, nsite, 4)
@@ -793,6 +796,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
     mocap_bodyid: id of body for mocap                       (nmocap,)
@@ -915,7 +920,9 @@ class Model:
   geom_margin: wp.array2d(dtype=float)
   geom_gap: wp.array2d(dtype=float)
   geom_rgba: wp.array2d(dtype=wp.vec4)
+  site_type: wp.array(dtype=int)
   site_bodyid: wp.array(dtype=int)
+  site_size: wp.array(dtype=wp.vec3)
   site_pos: wp.array2d(dtype=wp.vec3)
   site_quat: wp.array2d(dtype=wp.quat)
   cam_mode: wp.array(dtype=int)
@@ -1033,6 +1040,7 @@ class Model:
   sensor_pos_adr: wp.array(dtype=int)  # warp only
   sensor_vel_adr: wp.array(dtype=int)  # warp only
   sensor_acc_adr: wp.array(dtype=int)  # warp only
+  sensor_touch_adr: wp.array(dtype=int)  # warp only
   sensor_subtree_vel: bool  # warp only
   sensor_rne_postconstraint: bool  # warp only
   mocap_bodyid: wp.array(dtype=int)  # warp only