From cf02171bf8d9d90603de193fd59491225f240dc4 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Fri, 30 Aug 2024 08:38:51 -0700 Subject: [PATCH] Add velocimeter and gyro sensors to MJX. PiperOrigin-RevId: 669337697 Change-Id: I5e7f245b96eb334ffddd01757a9f631d9ad862cd --- doc/changelog.rst | 11 ++--------- mjx/mujoco/mjx/_src/sensor.py | 18 +++++++++++++++++- mjx/mujoco/mjx/_src/types.py | 4 ++++ mjx/mujoco/mjx/test_data/sensor.xml | 6 ++++++ 4 files changed, 29 insertions(+), 10 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 7cc4f827..a1529bb9 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -16,10 +16,6 @@ General the :ref:`mjOption` struct. This is because the new convex collision detection pipeline (see below) does not use the MPR algorithm. The semantics of these options remain identical. -- Added a new way of defining :ref:`connect` equality constraints, using two sites rather than bodies. - The new semantic is useful when the assumption that the constraint is satisfied in the base configuration does not - hold. In this case the sites will "snap together" at the beginning of the simulation. Additionally, changing the site - positions in ``mjModel.site_pos`` at runtime can be used to modify the constraint. - Added the :ref:`nativeccd` flag. When this flag is enabled, general convex collision detection is handled natively, as opposed to using `libccd `__. This feature is in early stages of testing. @@ -28,7 +24,7 @@ General - When :ref:`attaching` sub-models, :ref:`keyframes` will now be correctly merged into the parent model, but only on the first attachment. - Added the :ref:`mjtSameFrame` enum which contains the possible frame alignments of bodies and their children. These - alignments are used for computation shortcuts in :ref:`mj_kinematics`. + alignments are used as shortcuts in :ref:`mj_kinematics`. MJX ^^^ @@ -36,13 +32,12 @@ MJX - Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``. -- Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. +- Added velocity-dependent sensors: ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. - 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``. - Added support for :ref:`implicitfast integration` for all cases except :doc:`fluid drag `. -- Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices. Bug fixes ^^^^^^^^^ @@ -50,8 +45,6 @@ Bug fixes contribution by :github:user:`michael-ahn`). - 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. -- Fixed a bug in tendon wrapping around spheres. Before this fix, tendons that wrapped around spheres with an - externally-placed :ref:`sidesite` could jump inside the sphere instead of wrapping around it. Python bindings ^^^^^^^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index d2ad1497..bb090496 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -266,7 +266,23 @@ def sensor_vel(m: Model, d: Data) -> Data: objid = m.sensor_objid[idx] adr = m.sensor_adr[idx] - if sensor_type == SensorType.JOINTVEL: + if sensor_type == SensorType.VELOCIMETER: + bodyid = m.site_bodyid[objid] + pos = d.site_xpos[objid] + rot = d.site_xmat[objid] + cvel = d.cvel[bodyid] + subtree_com = d.subtree_com[m.body_rootid[bodyid]] + sensor = jax.vmap( + lambda vec, dif, rot: rot.T @ (vec[3:] - jp.cross(dif, vec[:3])) + )(cvel, pos - subtree_com, rot).reshape(-1) + adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + elif sensor_type == SensorType.GYRO: + bodyid = m.site_bodyid[objid] + rot = d.site_xmat[objid] + ang = d.cvel[bodyid, :3] + sensor = jax.vmap(lambda ang, rot: rot.T @ ang)(ang, rot).reshape(-1) + adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + elif sensor_type == SensorType.JOINTVEL: sensor = d.qvel[m.jnt_dofadr[objid]] elif sensor_type == SensorType.ACTUATORVEL: sensor = d.actuator_velocity[objid] diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index f913914e..8f47d048 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -307,6 +307,8 @@ class SensorType(enum.IntEnum): FRAMEQUAT: frame orientation, represented as quaternion SUBTREECOM: subtree centor of mass CLOCK: simulation time + VELOCIMETER: 3D linear velocity, in local frame + GYRO: 3D angular velocity, in local frame JOINTVEL: joint velocity ACTUATORVEL: actuator velocity BALLANGVEL: ball joint angular velocity @@ -326,6 +328,8 @@ class SensorType(enum.IntEnum): FRAMEQUAT = mujoco.mjtSensor.mjSENS_FRAMEQUAT SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM CLOCK = mujoco.mjtSensor.mjSENS_CLOCK + VELOCIMETER = mujoco.mjtSensor.mjSENS_VELOCIMETER + GYRO = mujoco.mjtSensor.mjSENS_GYRO JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 5a636e31..c1231457 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -15,6 +15,8 @@ -subtreecom -clock * velocity-dependent sensors: +-velocimeter +-gyro -jointvel -actuatorvel -ballangvel @@ -82,6 +84,8 @@ + + @@ -96,6 +100,8 @@ + +