From 806b8c8e2e61fc9edfa69b4ed19ac2a2a4ca3d58 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 21 Aug 2024 03:45:40 -0700 Subject: [PATCH] Add support for actuatorfrc and jointactfrc sensors in MJX. PiperOrigin-RevId: 665804850 Change-Id: I99371fc4da055382f1564b2c67d25f7b30a0a0cf --- doc/changelog.rst | 13 +++++++------ mjx/mujoco/mjx/_src/sensor.py | 28 +++++++++++++++++++++++++++- mjx/mujoco/mjx/_src/types.py | 4 ++++ mjx/mujoco/mjx/test_data/sensor.xml | 6 ++++++ 4 files changed, 44 insertions(+), 7 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 9b47e771..695122ee 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -21,21 +21,22 @@ MJX ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. 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``. -10. Added support for :ref:`implicitfast integration` for all cases except +8. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. +9. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. +10. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +11. Added support for :ref:`implicitfast integration` for all cases except :doc:`fluid drag `. Bug fixes ^^^^^^^^^ -11. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +12. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -12. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +13. 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 ^^^^^^^^^^^^^^^ -13. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +14. 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 77f6a3aa..4e704472 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -263,4 +263,30 @@ def sensor_acc(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d - return d + stage_acc = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC + sensors, adrs = [], [] + + for sensor_type in set(m.sensor_type[stage_acc]): + idx = m.sensor_type == sensor_type + objid = m.sensor_objid[idx] + adr = m.sensor_adr[idx] + + if sensor_type == SensorType.ACTUATORFRC: + sensor = d.actuator_force[objid] + elif sensor_type == SensorType.JOINTACTFRC: + sensor = d.qfrc_actuator[m.jnt_dofadr[objid]] + else: + # TODO(taylorhowell): raise error after adding sensor check to io.py + 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) + ) + + return d.replace(sensordata=sensordata) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 4a14db17..fd3aa362 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -309,6 +309,8 @@ class SensorType(enum.IntEnum): JOINTVEL: joint velocity ACTUATORVEL: actuator velocity BALLANGVEL: ball joint angular velocity + ACTUATORFRC: scalar actuator force + JOINTACTFRC: scalar actuator force, measured at the joint """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION @@ -325,6 +327,8 @@ class SensorType(enum.IntEnum): JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL + ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC + JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC class ObjType(PyTreeNode): diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 04924952..80285609 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -18,6 +18,8 @@ -actuatorvel -ballangvel * acceleration/force-dependent sensors: +-actuatorfrc +-jointactfrc --> @@ -81,14 +83,17 @@ + + + @@ -97,6 +102,7 @@ +