From 20296466b102cfd898cc2f7fda2fd923ffcd5434 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 5 Sep 2024 06:20:39 -0700 Subject: [PATCH] Add subtreelinvel and subtreeangmom sensors to MJX. PiperOrigin-RevId: 671347528 Change-Id: I8a32590c650ca6d0bb6df43d0347f707ffc05f12 --- doc/changelog.rst | 2 +- mjx/mujoco/mjx/_src/sensor.py | 17 ++++++++++++++--- mjx/mujoco/mjx/_src/types.py | 4 ++++ mjx/mujoco/mjx/test_data/sensor/sensor.xml | 6 ++++++ 4 files changed, 25 insertions(+), 4 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 950a4262..d045ca61 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -53,7 +53,7 @@ MJX ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``. - Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, - ``FRAMELINVEL``, ``FRAMEANGVEL``. + ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``. - Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. - Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. - Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index 05bd721d..7eb1b315 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -20,6 +20,7 @@ import mujoco # pylint: disable=g-importing-member from mujoco.mjx._src import math from mujoco.mjx._src import ray +from mujoco.mjx._src import smooth from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import Model @@ -296,9 +297,13 @@ def sensor_vel(m: Model, d: Data) -> Data: } stage_vel = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL - sensors, adrs = [], [] + sensor_types = set(m.sensor_type[stage_vel]) - for sensor_type in set(m.sensor_type[stage_vel]): + if sensor_types & {SensorType.SUBTREELINVEL, SensorType.SUBTREEANGMOM}: + d = smooth.subtree_vel(m, d) + + sensors, adrs = [], [] + for sensor_type in sensor_types: idx = m.sensor_type == sensor_type objid = m.sensor_objid[idx] adr = m.sensor_adr[idx] @@ -329,7 +334,7 @@ def sensor_vel(m: Model, d: Data) -> Data: jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None] sensor = d.qvel[jnt_dotadr] adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) - elif sensor_type in (SensorType.FRAMELINVEL, SensorType.FRAMEANGVEL): + elif sensor_type in {SensorType.FRAMELINVEL, SensorType.FRAMEANGVEL}: objtype = m.sensor_objtype[idx] reftype = m.sensor_reftype[idx] refid = m.sensor_refid[idx] @@ -385,6 +390,12 @@ def sensor_vel(m: Model, d: Data) -> Data: sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1)) adrs.append(adrt.reshape(-1)) continue # avoid adding to sensors/adrs list a second time + elif sensor_type == SensorType.SUBTREELINVEL: + sensor = d.subtree_linvel[objid] + adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + elif sensor_type == SensorType.SUBTREEANGMOM: + sensor = d.subtree_angmom[objid] + adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) else: # TODO(taylorhowell): raise error after adding sensor check to io.py continue # unsupported sensor type diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index ffe1ad5a..def0899b 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -316,6 +316,8 @@ class SensorType(enum.IntEnum): BALLANGVEL: ball joint angular velocity FRAMELINVEL: 3D linear velocity FRAMEANGVEL: 3D angular velocity + SUBTREELINVEL: subtree linear velocity + SUBTREEANGMOM: subtree angular momentum ACTUATORFRC: scalar actuator force JOINTACTFRC: scalar actuator force, measured at the joint """ @@ -339,6 +341,8 @@ class SensorType(enum.IntEnum): BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL FRAMELINVEL = mujoco.mjtSensor.mjSENS_FRAMELINVEL FRAMEANGVEL = mujoco.mjtSensor.mjSENS_FRAMEANGVEL + SUBTREELINVEL = mujoco.mjtSensor.mjSENS_SUBTREELINVEL + SUBTREEANGMOM = mujoco.mjtSensor.mjSENS_SUBTREEANGMOM ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC diff --git a/mjx/mujoco/mjx/test_data/sensor/sensor.xml b/mjx/mujoco/mjx/test_data/sensor/sensor.xml index 350c7bd3..0c5ca76b 100644 --- a/mjx/mujoco/mjx/test_data/sensor/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor/sensor.xml @@ -22,6 +22,8 @@ -ballangvel -framelinvel -frameangvel +-subtreelinvel +-subtreeangmom * acceleration/force-dependent sensors: -actuatorfrc -jointactfrc @@ -98,6 +100,7 @@ + @@ -114,15 +117,18 @@ + + +