From 7bc7c8a576f89d94924710c88cc470a8a1108ff4 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Sat, 3 May 2025 19:29:38 +0100 Subject: [PATCH] add touch sensor --- mujoco_warp/_src/io.py | 10 ++++- mujoco_warp/_src/sensor.py | 70 +++++++++++++++++++++++++++++++-- mujoco_warp/_src/sensor_test.py | 44 +++++++++++++++++++++ mujoco_warp/_src/types.py | 5 +++ 4 files changed, 124 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 87c1bdf3..76b5e59d 100644 --- a/mujoco_warp/_src/io.py +++ b/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, ) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index f4a10c4e..c5d7fd19 100644 --- a/mujoco_warp/_src/sensor.py +++ b/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]) diff --git a/mujoco_warp/_src/sensor_test.py b/mujoco_warp/_src/sensor_test.py index 50797234..e0344eac 100644 --- a/mujoco_warp/_src/sensor_test.py +++ b/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=""" + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + """, + 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) diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 1bf75b81..e79fb6a3 100644 --- a/mujoco_warp/_src/types.py +++ b/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