From 096853e192077dcec3e467e308846c8eaec36873 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Tue, 15 Oct 2024 06:37:00 -0700 Subject: [PATCH] Add JOINTINPARENT actuator transmission type to MJX. PiperOrigin-RevId: 686083761 Change-Id: I5159d050812d32e15ced06df132d3e434fced814 --- doc/changelog.rst | 1 + doc/mjx.rst | 4 ++-- mjx/mujoco/mjx/_src/passive_test.py | 4 +++- mjx/mujoco/mjx/_src/scan.py | 4 ++-- mjx/mujoco/mjx/_src/smooth.py | 14 +++++++++++--- mjx/mujoco/mjx/_src/types.py | 5 ++++- mjx/mujoco/mjx/test_data/pendula.xml | 4 +++- 7 files changed, 26 insertions(+), 10 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 29e90f2c..2ed3f095 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -26,6 +26,7 @@ MJX - Fixed a bug with frictionloss constraints. - Added ``TENDONPOS`` and ``TENDONVEL`` sensors. - Fixed a bug with the computation of tangential contact forces in ``_decode_pyramid``. +- Added ``JOINTINPARENT`` actuator transmission type. Python bindings ^^^^^^^^^^^^^^^ diff --git a/doc/mjx.rst b/doc/mjx.rst index 7bbf1066..9c982edc 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -188,7 +188,7 @@ The following features are **fully supported** in MJX: * - :ref:`Joint ` - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE`` * - :ref:`Transmission ` - - ``TRN_JOINT``, ``TRN_SITE``, ``TRN_TENDON`` + - ``JOINT``, ``JOINTINPARENT``, ``SITE``, ``TENDON`` * - :ref:`Actuator Dynamics ` - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT`` * - :ref:`Actuator Gain ` @@ -270,7 +270,7 @@ The following features are **unsupported**: * - :ref:`margin` and :ref:`gap` - Unimplemented for collisions with ``Mesh`` :ref:`Geom `. * - :ref:`Transmission ` - - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_BODY`` + - ``SLIDERCRANK``, ``BODY`` * - :ref:`Actuator Dynamics ` - ``USER`` * - :ref:`Actuator Gain ` diff --git a/mjx/mujoco/mjx/_src/passive_test.py b/mjx/mujoco/mjx/_src/passive_test.py index 30bc8c66..4ac3cd40 100644 --- a/mjx/mujoco/mjx/_src/passive_test.py +++ b/mjx/mujoco/mjx/_src/passive_test.py @@ -42,7 +42,9 @@ class PassiveTest(absltest.TestCase): m = test_util.load_test_file('pendula.xml') d = mujoco.MjData(m) # give the system a little kick to ensure we have non-identity rotations - d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4, 0.5, -0.6, 0.1, -0.2]) + d.ctrl = np.array( + [0.1, -0.1, 0.2, 0.3, -0.4, 0.5, -0.6, 0.1, -0.2, 0.1, 0.2] + ) mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero mujoco.mj_forward(m, d) mx = mjx.put_model(m) diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py index a0cf3037..cbc9bd0e 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/mjx/mujoco/mjx/_src/scan.py @@ -240,7 +240,7 @@ def flat( 'a': m.actuator_actadr[i], 'j': ( m.actuator_trnid[i, 0] - if m.actuator_trntype[i] == TrnType.JOINT + if m.actuator_trntype[i] in (TrnType.JOINT, TrnType.JOINTINPARENT) else -1 ), 's': ( @@ -250,7 +250,7 @@ def flat( ), } v, q = np.array([-1]), np.array([-1]) - if m.actuator_trntype[i] == TrnType.JOINT: + if m.actuator_trntype[i] in (TrnType.JOINT, TrnType.JOINTINPARENT): # v/q are associated with the joint transmissions only v = np.nonzero(m.dof_jntid == typ_ids['j'])[0] q = np.nonzero(_q_jointid(m) == typ_ids['j'])[0] diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 6eb8491b..3a168221 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -1022,15 +1022,23 @@ def transmission(m: Model, d: Data) -> Data: site_xmat, site_quat, ): - if trntype == TrnType.JOINT: + if trntype in (TrnType.JOINT, TrnType.JOINTINPARENT): if jnt_typ == JointType.FREE: length = jp.zeros(1) moment = gear + if trntype == TrnType.JOINTINPARENT: + quat_neg = math.quat_inv(qpos[3:]) + gearaxis = math.rotate(gear[3:], quat_neg) + moment = moment.at[3:].set(gearaxis) m_j = m_j + jp.arange(6) elif jnt_typ == JointType.BALL: axis, angle = math.quat_to_axis_angle(qpos) - length = jp.dot(axis * angle, gear[:3])[None] - moment = gear[:3] + gearaxis = gear[:3] + if trntype == TrnType.JOINTINPARENT: + quat_neg = math.quat_inv(qpos) + gearaxis = math.rotate(gear[:3], quat_neg) + length = jp.dot(axis * angle, gearaxis)[None] + moment = gearaxis m_j = m_j + jp.arange(3) elif jnt_typ in (JointType.SLIDE, JointType.HINGE): length = qpos * gear[0] diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index f886d3a3..15eaa8c3 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -215,12 +215,15 @@ class TrnType(enum.IntEnum): Members: JOINT: force on joint + JOINTINPARENT: force on joint, expressed in parent frame + TENDON: force on tendon SITE: force on site """ JOINT = mujoco.mjtTrn.mjTRN_JOINT + JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT SITE = mujoco.mjtTrn.mjTRN_SITE TENDON = mujoco.mjtTrn.mjTRN_TENDON - # unsupported: JOINTINPARENT, SLIDERCRANK, BODY + # unsupported: SLIDERCRANK, BODY class DynType(enum.IntEnum): diff --git a/mjx/mujoco/mjx/test_data/pendula.xml b/mjx/mujoco/mjx/test_data/pendula.xml index d183e5e9..6902b2f9 100644 --- a/mjx/mujoco/mjx/test_data/pendula.xml +++ b/mjx/mujoco/mjx/test_data/pendula.xml @@ -99,7 +99,7 @@ - + @@ -154,5 +154,7 @@ + +