From 6241814ad79d766c6ff51067ba000f1196d09a08 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Sun, 18 Aug 2024 08:46:45 -0700 Subject: [PATCH] Add velocity-dependent sensors to MJX. Includes: joint velocity, actuator velocity, and ball joint angular velocity. PiperOrigin-RevId: 664461284 Change-Id: I5063f662d707fbe342493f9cd7993f009fac9c76 --- doc/changelog.rst | 11 ++++---- mjx/mujoco/mjx/_src/sensor.py | 39 ++++++++++++++++++++++++----- mjx/mujoco/mjx/_src/types.py | 6 +++++ mjx/mujoco/mjx/test_data/sensor.xml | 9 +++++++ 4 files changed, 54 insertions(+), 11 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 4ddf86f1..87cbebea 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -20,19 +20,20 @@ MJX 6. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. -7. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. -8. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. +8. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. +9. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. Bug fixes ^^^^^^^^^ -9. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +10. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -10. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +11. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit integrators, wrong derivatives would be computed. Python bindings ^^^^^^^^^^^^^^^ -11. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +12. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index f8792764..77f6a3aa 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -35,11 +35,6 @@ def sensor_pos(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d - # no position-dependent sensors - stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS - if sum(stage_pos) == 0: - return d - # position and orientation by object type objtype_data = { ObjType.UNKNOWN: ( @@ -60,6 +55,7 @@ def sensor_pos(m: Model, d: Data) -> Data: SensorType.FRAMEZAXIS: 2, } + stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS sensors, adrs = [], [] for sensor_type in set(m.sensor_type[stage_pos]): @@ -206,6 +202,7 @@ def sensor_pos(m: Model, d: Data) -> Data: elif sensor_type == SensorType.CLOCK: sensor = jp.repeat(d.time, sum(idx)) else: + # TODO(taylorhowell): raise error after adding sensor check to io.py continue # unsupported sensor type sensors.append(sensor) @@ -227,7 +224,37 @@ def sensor_vel(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d - return d + stage_vel = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL + sensors, adrs = [], [] + + for sensor_type in set(m.sensor_type[stage_vel]): + idx = m.sensor_type == sensor_type + objid = m.sensor_objid[idx] + adr = m.sensor_adr[idx] + + if sensor_type == SensorType.JOINTVEL: + sensor = d.qvel[m.jnt_dofadr[objid]] + elif sensor_type == SensorType.ACTUATORVEL: + sensor = d.actuator_velocity[objid] + elif sensor_type == SensorType.BALLANGVEL: + jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None] + sensor = d.qvel[jnt_dotadr].reshape(-1) + adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + else: + # TODO(taylorhowell): raise error after adding sensor check to io.py + continue # unsupported sensor typ + + sensors.append(sensor) + adrs.append(adr) + + if not adrs: + return d + + sensordata = d.sensordata.at[np.concatenate(adrs)].set( + jp.concatenate(sensors) + ) + + return d.replace(sensordata=sensordata) def sensor_acc(m: Model, d: Data) -> Data: diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index de53a7a1..15e8bb18 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -304,6 +304,9 @@ class SensorType(enum.IntEnum): FRAMEZAXIS: frame z-axis SUBTREECOM: subtree centor of mass CLOCK: simulation time + JOINTVEL: joint velocity + ACTUATORVEL: actuator velocity + BALLANGVEL: ball joint angular velocity """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION @@ -317,6 +320,9 @@ class SensorType(enum.IntEnum): FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM CLOCK = mujoco.mjtSensor.mjSENS_CLOCK + JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL + ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL + BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL class ObjType(PyTreeNode): diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 72b21e3f..04924952 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -14,6 +14,9 @@ -subtreecom -clock * velocity-dependent sensors: +-jointvel +-actuatorvel +-ballangvel * acceleration/force-dependent sensors: --> @@ -77,8 +80,11 @@ + + + @@ -87,12 +93,15 @@ + + +