Skip to content

Commit 703f3e0

Browse files
authored
Merge pull request #213 from thowell/sensor_touch
add touch sensor
2 parents d158242 + 7982dd8 commit 703f3e0

File tree

5 files changed

+235
-38
lines changed

5 files changed

+235
-38
lines changed

mujoco_warp/_src/io.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -405,7 +405,9 @@ def create_nmodel_batched_array(mjm_array, dtype):
405405
geom_margin=create_nmodel_batched_array(mjm.geom_margin, dtype=float),
406406
geom_gap=create_nmodel_batched_array(mjm.geom_gap, dtype=float),
407407
geom_rgba=create_nmodel_batched_array(mjm.geom_rgba, dtype=wp.vec4),
408+
site_type=wp.array(mjm.site_type, dtype=int),
408409
site_bodyid=wp.array(mjm.site_bodyid, dtype=int),
410+
site_size=wp.array(mjm.site_size, dtype=wp.vec3),
409411
site_pos=create_nmodel_batched_array(mjm.site_pos, dtype=wp.vec3),
410412
site_quat=create_nmodel_batched_array(mjm.site_quat, dtype=wp.quat),
411413
cam_mode=wp.array(mjm.cam_mode, dtype=int),
@@ -534,7 +536,11 @@ def create_nmodel_batched_array(mjm_array, dtype):
534536
dtype=int,
535537
),
536538
sensor_acc_adr=wp.array(
537-
np.nonzero(mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC)[0],
539+
np.nonzero((mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC) & (mjm.sensor_type != mujoco.mjtSensor.mjSENS_TOUCH))[0],
540+
dtype=int,
541+
),
542+
sensor_touch_adr=wp.array(
543+
np.nonzero(mjm.sensor_type == mujoco.mjtSensor.mjSENS_TOUCH)[0],
538544
dtype=int,
539545
),
540546
sensor_subtree_vel=np.isin(

mujoco_warp/_src/ray.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
# limitations under the License.
1414
# ==============================================================================
1515

16+
from typing import Tuple
17+
1618
import warp as wp
1719

1820
from .types import MJ_MINVAL
@@ -361,7 +363,7 @@ def _ray_map(
361363
mat: wp.mat33,
362364
pnt: wp.vec3,
363365
vec: wp.vec3,
364-
) -> any:
366+
) -> Tuple[wp.vec3, wp.vec3]:
365367
"""Maps ray to local geom frame coordinates.
366368
367369
Args:
@@ -384,7 +386,7 @@ def _ray_map(
384386

385387

386388
@wp.func
387-
def _ray_geom(
389+
def ray_geom(
388390
# In:
389391
pos: wp.vec3, # Position of geom frame
390392
mat: wp.mat33, # Orientation of geom frame

mujoco_warp/_src/sensor.py

Lines changed: 172 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@
1818
import warp as wp
1919

2020
from . import math
21+
from . import ray
2122
from . import smooth
2223
from .types import MJ_MINVAL
24+
from .types import ConeType
2325
from .types import Data
2426
from .types import DataType
2527
from .types import DisableBit
@@ -1273,44 +1275,179 @@ def _sensor_acc(
12731275
_write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out)
12741276

12751277

1278+
@wp.kernel
1279+
def _sensor_touch_zero(
1280+
# Model:
1281+
sensor_adr: wp.array(dtype=int),
1282+
sensor_touch_adr: wp.array(dtype=int),
1283+
# Data out:
1284+
sensordata_out: wp.array2d(dtype=float),
1285+
):
1286+
worldid, sensortouchadrid = wp.tid()
1287+
sensorid = sensor_touch_adr[sensortouchadrid]
1288+
adr = sensor_adr[sensorid]
1289+
sensordata_out[worldid, adr] = 0.0
1290+
1291+
1292+
@wp.kernel
1293+
def _sensor_touch(
1294+
# Model:
1295+
opt_cone: int,
1296+
geom_bodyid: wp.array(dtype=int),
1297+
site_type: wp.array(dtype=int),
1298+
site_bodyid: wp.array(dtype=int),
1299+
site_size: wp.array(dtype=wp.vec3),
1300+
sensor_objid: wp.array(dtype=int),
1301+
sensor_adr: wp.array(dtype=int),
1302+
sensor_touch_adr: wp.array(dtype=int),
1303+
# Data in:
1304+
ncon_in: wp.array(dtype=int),
1305+
site_xpos_in: wp.array2d(dtype=wp.vec3),
1306+
site_xmat_in: wp.array2d(dtype=wp.mat33),
1307+
contact_pos_in: wp.array(dtype=wp.vec3),
1308+
contact_frame_in: wp.array(dtype=wp.mat33),
1309+
contact_dim_in: wp.array(dtype=int),
1310+
contact_geom_in: wp.array(dtype=wp.vec2i),
1311+
contact_efc_address_in: wp.array2d(dtype=int),
1312+
contact_worldid_in: wp.array(dtype=int),
1313+
efc_force_in: wp.array(dtype=float),
1314+
# Data out:
1315+
sensordata_out: wp.array2d(dtype=float),
1316+
):
1317+
conid, sensortouchadrid = wp.tid()
1318+
1319+
if conid > ncon_in[0]:
1320+
return
1321+
1322+
sensorid = sensor_touch_adr[sensortouchadrid]
1323+
1324+
objid = sensor_objid[sensorid]
1325+
bodyid = site_bodyid[objid]
1326+
1327+
# find contact in sensor zone, add normal force
1328+
1329+
# contacting bodies
1330+
geom = contact_geom_in[conid]
1331+
conbody = wp.vec2i(geom_bodyid[geom[0]], geom_bodyid[geom[1]])
1332+
1333+
# select contacts involving sensorized body
1334+
efc_address0 = contact_efc_address_in[conid, 0]
1335+
if efc_address0 >= 0 and (bodyid == conbody[0] or bodyid == conbody[1]):
1336+
# get contact normal force
1337+
normalforce = efc_force_in[efc_address0]
1338+
1339+
if opt_cone == int(ConeType.PYRAMIDAL.value):
1340+
dim = contact_dim_in[conid]
1341+
for i in range(1, 2 * (dim - 1)):
1342+
normalforce += efc_force_in[contact_efc_address_in[conid, i]]
1343+
1344+
if normalforce <= 0.0:
1345+
return
1346+
1347+
# convert contact normal force to global frame, normalize
1348+
frame = contact_frame_in[conid]
1349+
conray = wp.vec3(frame[0, 0], frame[0, 1], frame[0, 2]) * normalforce
1350+
conray, _ = math.normalize_with_norm(conray)
1351+
1352+
# flip ray direction if sensor is on body2
1353+
if bodyid == conbody[1]:
1354+
conray = -conray
1355+
1356+
# add if ray-zone intersection (always true when contact.pos inside zone)
1357+
worldid = contact_worldid_in[conid]
1358+
if (
1359+
ray.ray_geom(
1360+
site_xpos_in[worldid, objid],
1361+
site_xmat_in[worldid, objid],
1362+
site_size[objid],
1363+
contact_pos_in[conid],
1364+
conray,
1365+
site_type[objid],
1366+
)
1367+
>= 0.0
1368+
):
1369+
adr = sensor_adr[sensorid]
1370+
wp.atomic_add(sensordata_out[worldid], adr, normalforce)
1371+
1372+
12761373
@event_scope
12771374
def sensor_acc(m: Model, d: Data):
12781375
"""Compute acceleration-dependent sensor values."""
1279-
1280-
if (m.sensor_acc_adr.size == 0) or (m.opt.disableflags & DisableBit.SENSOR):
1376+
if m.opt.disableflags & DisableBit.SENSOR:
12811377
return
12821378

1283-
if m.sensor_rne_postconstraint:
1284-
smooth.rne_postconstraint(m, d)
1379+
if m.sensor_touch_adr.size > 0:
1380+
wp.launch(
1381+
_sensor_touch_zero,
1382+
dim=(d.nworld, m.sensor_touch_adr.size),
1383+
inputs=[
1384+
m.sensor_adr,
1385+
m.sensor_touch_adr,
1386+
],
1387+
outputs=[
1388+
d.sensordata,
1389+
],
1390+
)
1391+
wp.launch(
1392+
_sensor_touch,
1393+
dim=(d.nconmax, m.sensor_touch_adr.size),
1394+
inputs=[
1395+
m.opt.cone,
1396+
m.geom_bodyid,
1397+
m.site_type,
1398+
m.site_bodyid,
1399+
m.site_size,
1400+
m.sensor_objid,
1401+
m.sensor_adr,
1402+
m.sensor_touch_adr,
1403+
d.ncon,
1404+
d.site_xpos,
1405+
d.site_xmat,
1406+
d.contact.pos,
1407+
d.contact.frame,
1408+
d.contact.dim,
1409+
d.contact.geom,
1410+
d.contact.efc_address,
1411+
d.contact.worldid,
1412+
d.efc.force,
1413+
],
1414+
outputs=[
1415+
d.sensordata,
1416+
],
1417+
)
12851418

1286-
wp.launch(
1287-
_sensor_acc,
1288-
dim=(d.nworld, m.sensor_acc_adr.size),
1289-
inputs=[
1290-
m.body_rootid,
1291-
m.jnt_dofadr,
1292-
m.geom_bodyid,
1293-
m.site_bodyid,
1294-
m.cam_bodyid,
1295-
m.sensor_type,
1296-
m.sensor_datatype,
1297-
m.sensor_objtype,
1298-
m.sensor_objid,
1299-
m.sensor_adr,
1300-
m.sensor_cutoff,
1301-
m.sensor_acc_adr,
1302-
d.xpos,
1303-
d.xipos,
1304-
d.geom_xpos,
1305-
d.site_xpos,
1306-
d.site_xmat,
1307-
d.cam_xpos,
1308-
d.subtree_com,
1309-
d.cvel,
1310-
d.actuator_force,
1311-
d.qfrc_actuator,
1312-
d.cacc,
1313-
d.cfrc_int,
1314-
],
1315-
outputs=[d.sensordata],
1316-
)
1419+
if m.sensor_acc_adr.size > 0:
1420+
if m.sensor_rne_postconstraint:
1421+
smooth.rne_postconstraint(m, d)
1422+
1423+
wp.launch(
1424+
_sensor_acc,
1425+
dim=(d.nworld, m.sensor_acc_adr.size),
1426+
inputs=[
1427+
m.body_rootid,
1428+
m.jnt_dofadr,
1429+
m.geom_bodyid,
1430+
m.site_bodyid,
1431+
m.cam_bodyid,
1432+
m.sensor_type,
1433+
m.sensor_datatype,
1434+
m.sensor_objtype,
1435+
m.sensor_objid,
1436+
m.sensor_adr,
1437+
m.sensor_cutoff,
1438+
m.sensor_acc_adr,
1439+
d.xpos,
1440+
d.xipos,
1441+
d.geom_xpos,
1442+
d.site_xpos,
1443+
d.site_xmat,
1444+
d.cam_xpos,
1445+
d.subtree_com,
1446+
d.cvel,
1447+
d.actuator_force,
1448+
d.qfrc_actuator,
1449+
d.cacc,
1450+
d.cfrc_int,
1451+
],
1452+
outputs=[d.sensordata],
1453+
)

mujoco_warp/_src/sensor_test.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,50 @@ def test_sensor(self):
233233

234234
_assert_eq(d.sensordata.numpy()[0], mjd.sensordata, "sensordata")
235235

236+
def test_touch_sensor(self):
237+
"""Test touch sensor."""
238+
for keyframe in range(2):
239+
_, mjd, m, d = test_util.fixture(
240+
xml="""
241+
<mujoco>
242+
<worldbody>
243+
<geom type="plane" size="10 10 .001"/>
244+
<body pos="0 0 .25">
245+
<geom type="sphere" size="0.1" pos=".1 0 0"/>
246+
<geom type="sphere" size="0.1" pos="-.1 0 0"/>
247+
<geom type="sphere" size="0.1" pos="0 0 .11"/>
248+
<geom type="sphere" size="0.1" pos="0 0 10"/>
249+
<site name="site_sphere" type="sphere" size=".2"/>
250+
<site name="site_capsule" type="capsule" size=".2 .2"/>
251+
<site name="site_ellipsoid" type="ellipsoid" size=".2 .2 .2"/>
252+
<site name="site_cylinder" type="cylinder" size=".2 .2"/>
253+
<site name="site_box" type="box" size=".2 .2 .2"/>
254+
<freejoint/>
255+
</body>
256+
</worldbody>
257+
<sensor>
258+
<touch site="site_sphere"/>
259+
<touch site="site_capsule"/>
260+
<touch site="site_ellipsoid"/>
261+
<touch site="site_cylinder"/>
262+
<touch site="site_box"/>
263+
</sensor>
264+
<keyframe>
265+
<key qpos="0 0 10 1 0 0 0"/>
266+
<key qpos="0 0 .05 1 0 0 0"/>
267+
<key qpos="0 0 0 1 0 0 0"/>
268+
</keyframe>
269+
</mujoco>
270+
""",
271+
keyframe=keyframe,
272+
)
273+
274+
d.sensordata.zero_()
275+
276+
mjwarp.sensor_acc(m, d)
277+
278+
_assert_eq(d.sensordata.numpy()[0], mjd.sensordata, "sensordata")
279+
236280
def test_tendon_sensor(self):
237281
"""Test tendon sensors."""
238282
_, mjd, m, d = test_util.fixture("tendon/fixed.xml", keyframe=0, sparse=False)

mujoco_warp/_src/types.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,6 +262,7 @@ class SensorType(enum.IntEnum):
262262
FRAMEANGVEL: 3D angular velocity
263263
SUBTREELINVEL: subtree linear velocity
264264
SUBTREEANGMOM: subtree angular momentum
265+
TOUCH: scalar contact normal forces summed over sensor zone
265266
ACCELEROMETER: accelerometer
266267
FORCE: force
267268
TORQUE: torque
@@ -293,6 +294,7 @@ class SensorType(enum.IntEnum):
293294
FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL
294295
SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL
295296
SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM
297+
TOUCH = mujoco.mjtSensor.mjSENS_TOUCH
296298
ACCELEROMETER = mujoco.mjtSensor.mjSENS_ACCELEROMETER
297299
FORCE = mujoco.mjtSensor.mjSENS_FORCE
298300
TORQUE = mujoco.mjtSensor.mjSENS_TORQUE
@@ -687,6 +689,7 @@ class Model:
687689
geom_margin: detect contact if dist<margin (nworld, ngeom,)
688690
geom_gap: include in solver if dist<margin-gap (nworld, ngeom,)
689691
geom_rgba: rgba when material is omitted (nworld, ngeom, 4)
692+
site_type: geom type for rendering (mjtGeom) (nsite,)
690693
site_bodyid: id of site's body (nsite,)
691694
site_pos: local position offset rel. to body (nworld, nsite, 3)
692695
site_quat: local orientation offset rel. to body (nworld, nsite, 4)
@@ -793,6 +796,8 @@ class Model:
793796
sensor_pos_adr: addresses for position sensors (<=nsensor,)
794797
sensor_vel_adr: addresses for velocity sensors (<=nsensor,)
795798
sensor_acc_adr: addresses for acceleration sensors (<=nsensor,)
799+
(excluding touch sensors)
800+
sensor_touch_adr: addresses for touch sensors (<=nsensor,)
796801
sensor_subtree_vel: evaluate subtree_vel
797802
sensor_rne_postconstraint: evaluate rne_postconstraint
798803
mocap_bodyid: id of body for mocap (nmocap,)
@@ -915,7 +920,9 @@ class Model:
915920
geom_margin: wp.array2d(dtype=float)
916921
geom_gap: wp.array2d(dtype=float)
917922
geom_rgba: wp.array2d(dtype=wp.vec4)
923+
site_type: wp.array(dtype=int)
918924
site_bodyid: wp.array(dtype=int)
925+
site_size: wp.array(dtype=wp.vec3)
919926
site_pos: wp.array2d(dtype=wp.vec3)
920927
site_quat: wp.array2d(dtype=wp.quat)
921928
cam_mode: wp.array(dtype=int)
@@ -1033,6 +1040,7 @@ class Model:
10331040
sensor_pos_adr: wp.array(dtype=int) # warp only
10341041
sensor_vel_adr: wp.array(dtype=int) # warp only
10351042
sensor_acc_adr: wp.array(dtype=int) # warp only
1043+
sensor_touch_adr: wp.array(dtype=int) # warp only
10361044
sensor_subtree_vel: bool # warp only
10371045
sensor_rne_postconstraint: bool # warp only
10381046
mocap_bodyid: wp.array(dtype=int) # warp only

0 commit comments

Comments
 (0)