diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 9f095406..8dea42bb 100644 --- a/mujoco_warp/_src/io.py +++ b/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( diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index f877d1fd..e2aad274 100644 --- a/mujoco_warp/_src/ray.py +++ b/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 diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 2a5f43c0..2003ed9f 100644 --- a/mujoco_warp/_src/sensor.py +++ b/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], + ) diff --git a/mujoco_warp/_src/sensor_test.py b/mujoco_warp/_src/sensor_test.py index 2e7514d8..82938c92 100644 --- a/mujoco_warp/_src/sensor_test.py +++ b/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=""" + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + """, + 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 3ad48dd5..2c6d1acb 100644 --- a/mujoco_warp/_src/types.py +++ b/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