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

Modify MJX sensor_pos to handle unsupported sensors.

Browse Source 
PiperOrigin-RevId: 663768186
Change-Id: I2e7600e77855a96a47f012fb64b6cc43abbf6f64
This commit is contained in:
Taylor Howell
2024-08-16 10:09:42 -07:00
committed by Copybara-Service
parent 4d9305fbb9
commit ce9c0ea351
3 changed files with 34 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/sensor.py
+8 -2
View File
@@ -31,7 +31,8 @@ def sensor_pos(m: Model, d: Data) -> Data:
"""Compute position-dependent sensors values."""
# no position-dependent sensors
if sum(m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS) == 0:
stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS
if sum(stage_pos) == 0:
return d
# position and orientation by object type
@@ -56,7 +57,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
sensors, adrs = [], []
for sensor_type in set(m.sensor_type):
for sensor_type in set(m.sensor_type[stage_pos]):
idx = m.sensor_type == sensor_type
objid = m.sensor_objid[idx]
adr = m.sensor_adr[idx]
@@ -126,10 +127,15 @@ def sensor_pos(m: Model, d: Data) -> Data:
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
elif sensor_type == SensorType.CLOCK:
sensor = jp.repeat(d.time, sum(idx))
else:
continue # unsupported sensor type
sensors.append(sensor)
adrs.append(adr)
if not adrs:
return d
sensordata = d.sensordata.at[np.concatenate(adrs)].set(
jp.concatenate(sensors)
)
mjx/mujoco/mjx/_src/sensor_test.py
+7
View File
@@ -62,6 +62,13 @@ class SensorTest(parameterized.TestCase):
# sensor values
_assert_eq(d.sensordata, dx.sensordata, 'sensordata')
def test_unsupported_sensor(self):
"""Tests MJX sensor functions do not break for unsupported sensors."""
m = test_util.load_test_file('unsupported_sensor.xml')
mx = mjx.put_model(m)
dx = jax.jit(mjx.forward)(mx, mjx.put_data(m, mujoco.MjData(m)))
_assert_eq(np.zeros(m.nsensordata), dx.sensordata, 'sensordata')
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/test_data/unsupported_sensor.xml
+19
View File
@@ -0,0 +1,19 @@
<!-- For validating model with unsupported sensor -->
<mujoco model="unsupported_sensor">
<worldbody>
<body>
<joint type="hinge"/>
<geom name="geom0" size="1"/>
<site name="site"/>
</body>
<body>
<joint type="hinge"/>
<geom name="geom1" size="1"/>
</body>
</worldbody>
<sensor>
<distance name="distance" geom1="geom0" geom2="geom1"/>
<framelinvel name="framelinvel" objtype="site" objname="site"/>
<touch name="touch" site="site"/>
</sensor>
</mujoco>