From f9401776d49fecd3cd48962c30d42ecb09b6620c Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 16 Sep 2024 07:32:59 -0700 Subject: [PATCH] Add framelinacc and frameangacc sensors to MJX. PiperOrigin-RevId: 675141513 Change-Id: Ie38dba3e8be82e0554c7c9e47c6d07f3cdad29b8 --- doc/changelog.rst | 2 +- doc/mjx.rst | 3 +- mjx/mujoco/mjx/_src/sensor.py | 45 ++++++++++++++++++++++ mjx/mujoco/mjx/_src/types.py | 4 ++ mjx/mujoco/mjx/test_data/sensor/sensor.xml | 6 +++ 5 files changed, 58 insertions(+), 2 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 149571fe..f77c3195 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -64,7 +64,7 @@ MJX 14. Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``. 15. Added acceleration/force-dependent sensors: ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, ``ACTUATORFRC``, - ``JOINTACTFRC``. + ``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``. 16. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. 17. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. 18. Added support for :ref:`implicitfast integration` for all cases except diff --git a/doc/mjx.rst b/doc/mjx.rst index abfdb690..8322b62d 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -220,7 +220,8 @@ The following features are **fully supported** in MJX: - ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``, ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``, ``FRAMEANGVEL``, ``SUBTREELINVEL``, - ``SUBTREEANGMOM``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, ``ACTUATORFRC``, ``JOINTACTFRC``. + ``SUBTREEANGMOM``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``, ``ACTUATORFRC``, ``JOINTACTFRC``, ``FRAMELINACC``, + ``FRAMEANGACC``. The following features are **in development** and coming soon: diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index 07811d3c..19d7b8ef 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -419,6 +419,16 @@ def sensor_acc(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d + # position and bodyid by object type + objtype_data = { + ObjType.UNKNOWN: (np.zeros((1, 3)), np.arange(1)), + ObjType.BODY: (d.xipos, np.arange(m.nbody)), + ObjType.XBODY: (d.xpos, np.arange(m.nbody)), + ObjType.GEOM: (d.geom_xpos, m.geom_bodyid), + ObjType.SITE: (d.site_xpos, m.site_bodyid), + ObjType.CAMERA: (d.cam_xpos, m.cam_bodyid), + } + stage_acc = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC sensor_types = set(m.sensor_type[stage_acc]) @@ -478,6 +488,41 @@ def sensor_acc(m: Model, d: Data) -> Data: sensor = d.actuator_force[objid] elif sensor_type == SensorType.JOINTACTFRC: sensor = d.qfrc_actuator[m.jnt_dofadr[objid]] + elif sensor_type in (SensorType.FRAMELINACC, SensorType.FRAMEANGACC): + objtype = m.sensor_objtype[idx] + + for ot in set(objtype): + idxt = objtype == ot + objidt = objid[idxt] + pos, bodyid = objtype_data[ot] + pos = pos[objidt] + bodyid = bodyid[objidt] + cacc = d.cacc[bodyid] + + if sensor_type == SensorType.FRAMELINACC: + + @jax.vmap + def _framelinacc(cvel, cacc, offset): + ang = cvel[:3] + lin = cvel[3:] - jp.cross(offset, cvel[:3]) + acc = cacc[3:] - jp.cross(offset, cacc[:3]) + correction = jp.cross(ang, lin) + return acc + correction + + cvel = d.cvel[bodyid] + offset = pos - d.subtree_com[m.body_rootid[bodyid]] + + sensor = _framelinacc(cvel, cacc, offset).reshape(-1) + elif sensor_type == SensorType.FRAMEANGACC: + sensor = cacc[:, :3].reshape(-1) + else: + raise ValueError(f'Unknown sensor type: {sensor_type}') + + adrt = adr[idxt, None] + np.arange(3)[None] + + sensors.append(sensor.reshape(-1)) + adrs.append(adrt.reshape(-1)) + continue # avoid adding to sensors/adrs list a second time 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 17f5bf2e..17cd499b 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -325,6 +325,8 @@ class SensorType(enum.IntEnum): TORQUE: torque ACTUATORFRC: scalar actuator force JOINTACTFRC: scalar actuator force, measured at the joint + FRAMELINACC: 3D linear acceleration + FRAMEANGACC: 3D angular acceleration """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION @@ -353,6 +355,8 @@ class SensorType(enum.IntEnum): TORQUE = mujoco.mjtSensor.mjSENS_TORQUE ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC + FRAMELINACC = mujoco.mjtSensor.mjSENS_FRAMELINACC + FRAMEANGACC = mujoco.mjtSensor.mjSENS_FRAMEANGACC class ObjType(PyTreeNode): diff --git a/mjx/mujoco/mjx/test_data/sensor/sensor.xml b/mjx/mujoco/mjx/test_data/sensor/sensor.xml index 79129f2f..426d5be2 100644 --- a/mjx/mujoco/mjx/test_data/sensor/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor/sensor.xml @@ -30,6 +30,8 @@ -torque -actuatorfrc -jointactfrc +-framelinacc +-frameangacc --> @@ -144,6 +146,8 @@ + + @@ -160,9 +164,11 @@ + +