chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix bug in MJX contact sensor.

Browse Source 
PiperOrigin-RevId: 794011199
Change-Id: I709a8bd9d07bec980aa60abd0346d946b037cbbd
This commit is contained in:
Baruch Tabanpour
2025-08-12 02:22:53 -07:00
committed by Copybara-Service
parent 0a02e7b722
commit da9037e2b0
2 changed files with 6 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/sensor.py
+5 -8
View File
@@ -588,6 +588,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
size = nslotdata(dataspec) size = nslotdata(dataspec)
num = np.minimum(int(dim / size), ncon) num = np.minimum(int(dim / size), ncon)
nsensor = idx_ds.sum()
if objtype == ObjType.UNKNOWN and reftype == ObjType.UNKNOWN: if objtype == ObjType.UNKNOWN and reftype == ObjType.UNKNOWN:
# all contacts match # all contacts match
@@ -601,11 +602,9 @@ def sensor_acc(m: Model, d: Data) -> Data:
nfound = sum(is_contact) nfound = sum(is_contact)
# if duplicate sensor # if duplicate sensor
nsensor = idx_ds.sum() cid = jp.tile(cid, (nsensor,))
cid = np.tile(cid, (nsensor,)) nfound = jp.tile(nfound, (nsensor,))
match = np.tile(match[:num], (nsensor,)) flip = jp.ones((cid.size, 3))
nfound = np.tile(nfound, (nsensor,))
flip = np.ones((cid.size, 3))
elif objtype == ObjType.GEOM or reftype == ObjType.GEOM: elif objtype == ObjType.GEOM or reftype == ObjType.GEOM:
sensorid1 = objid[idx_ds] sensorid1 = objid[idx_ds]
sensorid2 = refid[idx_ds] sensorid2 = refid[idx_ds]
@@ -651,8 +650,6 @@ def sensor_acc(m: Model, d: Data) -> Data:
# number of contacts per sensor # number of contacts per sensor
nfound = (match * is_contact[None, :]).sum(axis=1) nfound = (match * is_contact[None, :]).sum(axis=1)
match = match[:, :num].reshape(-1)
# TODO(taylorhowell): matching criteria: body, subtree # TODO(taylorhowell): matching criteria: body, subtree
else: else:
@@ -683,7 +680,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
if dataspec & (1 << 6): # tangent if dataspec & (1 << 6): # tangent
slot.append(flip[:, 2, None] * d._impl.contact.frame[cid, 1]) slot.append(flip[:, 2, None] * d._impl.contact.frame[cid, 1])
found = is_contact[cid] & match found = jp.tile(jp.arange(num), nsensor) < jp.repeat(nfound, num)
sensors.append((found[:, None] * jp.hstack(slot)).reshape(-1)) sensors.append((found[:, None] * jp.hstack(slot)).reshape(-1))
adrs.append( adrs.append(
(adr[idx_ds][:, None] + np.arange(num * size)[None]).reshape(-1) (adr[idx_ds][:, None] + np.arange(num * size)[None]).reshape(-1)
mjx/mujoco/mjx/_src/sensor_test.py
+1 -4
View File
@@ -114,12 +114,11 @@ class SensorTest(parameterized.TestCase):
datas = list(datas) datas = list(datas)
contact_sensors = '' contact_sensors = ''
for num in [1, 2, 3, 4, 5]: for num in [1, 2, 4, 5]:
for data in datas: for data in datas:
data = ' '.join(data) data = ' '.join(data)
for reduce in ['mindist', 'maxforce']: for reduce in ['mindist', 'maxforce']:
for match in [ for match in [
'',
'', '',
'geom1="plane"', 'geom1="plane"',
'geom1="geom1"', 'geom1="geom1"',
@@ -130,9 +129,7 @@ class SensorTest(parameterized.TestCase):
'geom1="plane" geom2="geom1"', 'geom1="plane" geom2="geom1"',
'geom1="geom1" geom2="plane"', 'geom1="geom1" geom2="plane"',
'geom1="plane" geom2="sphere2"', 'geom1="plane" geom2="sphere2"',
'geom1="sphere2" geom2="plane"',
'geom1="geom1" geom2="sphere2"', 'geom1="geom1" geom2="sphere2"',
'geom1="sphere2" geom2="geom1"',
]: ]:
contact_sensors += ( contact_sensors += (
f'<contact {match} num="{num}" data="{data}"' f'<contact {match} num="{num}" data="{data}"'