Skip to content

add touch sensor #213

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
May 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion mujoco_warp/_src/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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(
Expand Down
6 changes: 4 additions & 2 deletions mujoco_warp/_src/ray.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@
# limitations under the License.
# ==============================================================================

from typing import Tuple

import warp as wp

from .types import MJ_MINVAL
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
207 changes: 172 additions & 35 deletions mujoco_warp/_src/sensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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],
)
44 changes: 44 additions & 0 deletions mujoco_warp/_src/sensor_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
8 changes: 8 additions & 0 deletions mujoco_warp/_src/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
Loading