From 0f381a9ebd6beca19cff5e42298b2b86e2940880 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 11 Nov 2024 07:45:11 -0800 Subject: [PATCH] Add muscle actuators to MJX. PiperOrigin-RevId: 695334356 Change-Id: I6590ed6bdd4a8adc53d5c4e2641f472f7439cedc --- doc/changelog.rst | 4 + doc/mjx.rst | 12 +- mjx/mujoco/mjx/_src/forward.py | 12 +- mjx/mujoco/mjx/_src/forward_test.py | 41 ++-- mjx/mujoco/mjx/_src/support.py | 147 +++++++++++++ mjx/mujoco/mjx/_src/support_test.py | 205 ++++++++++++++++++ mjx/mujoco/mjx/_src/types.py | 12 +- mjx/mujoco/mjx/test_data/actuator/arm21.xml | 37 ++++ mjx/mujoco/mjx/test_data/actuator/arm26.xml | 118 ++++++++++ .../test_data/actuator/general_dyntype.xml | 18 ++ 10 files changed, 566 insertions(+), 40 deletions(-) create mode 100644 mjx/mujoco/mjx/test_data/actuator/arm21.xml create mode 100644 mjx/mujoco/mjx/test_data/actuator/arm26.xml create mode 100644 mjx/mujoco/mjx/test_data/actuator/general_dyntype.xml diff --git a/doc/changelog.rst b/doc/changelog.rst index 3ccaccc1..87225bea 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -6,6 +6,10 @@ Changelog Upcoming version (not yet released) ----------------------------------- +MJX +^^^ +- Added muscle actuators. + Bug fixes ^^^^^^^^^ - Fixed :github:issue:`2212`, type error in ``mjx.get_data``. diff --git a/doc/mjx.rst b/doc/mjx.rst index 49435bc6..87dfd9e3 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -214,11 +214,11 @@ The following features are **fully supported** in MJX: * - :ref:`Transmission ` - ``JOINT``, ``JOINTINPARENT``, ``SITE``, ``TENDON`` * - :ref:`Actuator Dynamics ` - - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT`` + - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT``, ``MUSCLE`` * - :ref:`Actuator Gain ` - - ``FIXED``, ``AFFINE`` + - ``FIXED``, ``AFFINE``, ``MUSCLE`` * - :ref:`Actuator Bias ` - - ``NONE``, ``AFFINE`` + - ``NONE``, ``AFFINE``, ``MUSCLE`` * - :ref:`Tendon Wrapping ` - ``JOINT``, ``SITE``, ``PULLEY`` * - :ref:`Geom ` @@ -264,12 +264,6 @@ The following features are **in development** and coming soon: - ``IMPLICIT`` * - Dynamics - :ref:`Inverse ` - * - :ref:`Actuator Dynamics ` - - ``MUSCLE`` - * - :ref:`Actuator Gain ` - - ``MUSCLE`` - * - :ref:`Actuator Bias ` - - ``MUSCLE`` * - :ref:`Tendon Wrapping ` - ``SPHERE``, ``CYLINDER`` * - Fluid Model diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index c187d685..07f33a90 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -115,6 +115,8 @@ def fwd_actuation(m: Model, d: Data) -> Data: act_dot = ctrl elif dyn_typ in (DynType.FILTER, DynType.FILTEREXACT): act_dot = (ctrl - act) / jp.clip(dyn_prm[0], mujoco.mjMINVAL) + elif dyn_typ == DynType.MUSCLE: + act_dot = support.muscle_dynamics(ctrl, act, dyn_prm) else: raise NotImplementedError(f'dyntype {dyn_typ.name} not implemented.') return act_dot @@ -139,13 +141,15 @@ def fwd_actuation(m: Model, d: Data) -> Data: ctrl_act = jp.where(m.actuator_actadr == -1, ctrl, act_last_dim) def get_force(*args): - gain_t, gain_p, bias_t, bias_p, len_, vel, ctrl_act = args + gain_t, gain_p, bias_t, bias_p, len_, vel, ctrl_act, len_range, acc0 = args typ, prm = GainType(gain_t), gain_p if typ == GainType.FIXED: gain = prm[0] elif typ == GainType.AFFINE: gain = prm[0] + prm[1] * len_ + prm[2] * vel + elif typ == GainType.MUSCLE: + gain = support.muscle_gain(len_, vel, len_range, acc0, prm) else: raise RuntimeError(f'unrecognized gaintype {typ.name}.') @@ -153,13 +157,15 @@ def fwd_actuation(m: Model, d: Data) -> Data: bias = jp.array(0.0) if typ == BiasType.AFFINE: bias = prm[0] + prm[1] * len_ + prm[2] * vel + elif typ == BiasType.MUSCLE: + bias = support.muscle_bias(len_, len_range, acc0, prm) return gain * ctrl_act + bias force = scan.flat( m, get_force, - 'uuuuuuu', + 'uuuuuuuuu', 'u', m.actuator_gaintype, m.actuator_gainprm, @@ -168,6 +174,8 @@ def fwd_actuation(m: Model, d: Data) -> Data: d.actuator_length, d.actuator_velocity, ctrl_act, + jp.array(m.actuator_lengthrange), + jp.array(m.actuator_acc0), group_by='u', ) forcerange = jp.where( diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index 7c669062..ebe94ebe 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -15,6 +15,7 @@ """Tests for forward functions.""" from absl.testing import absltest +from absl.testing import parameterized import jax import mujoco from mujoco import mjx @@ -167,40 +168,28 @@ class ForwardTest(absltest.TestCase): np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep) -class ActuatorTest(absltest.TestCase): - _DYN_XML = """ - - - - - - - - - - - - - - - - - - - """ +class ActuatorTest(parameterized.TestCase): - def test_dyntype(self): - m = mujoco.MjModel.from_xml_string(self._DYN_XML) + @parameterized.parameters( + 'actuator/arm21.xml', + 'actuator/arm26.xml', + 'actuator/general_dyntype.xml', + ) + def test_actuator(self, fname): + m = test_util.load_test_file(fname) d = mujoco.MjData(m) - d.ctrl = np.array([1.5, 1.5, 1.5, 1.5]) - d.act = np.array([0.5, 0.5, 0.5]) - + mujoco.mj_step(m, d) + d.ctrl = 1.5 * np.random.random(m.nu) + d.act = 0.5 * np.random.random(m.na) mx = mjx.put_model(m) dx = mjx.put_data(m, d) mujoco.mj_fwdActuation(m, d) dx = jax.jit(mjx.fwd_actuation)(mx, dx) + _assert_attr_eq(d, dx, 'act_dot') + _assert_attr_eq(d, dx, 'qfrc_actuator') + _assert_attr_eq(d, dx, 'actuator_force') mujoco.mj_Euler(m, d) dx = jax.jit(mjx.euler)(mx, dx) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index f065d9d8..1f3c8137 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -560,3 +560,150 @@ def wrap( wpnt1 = jp.where(invalid, jp.zeros(3), wpnt1) return wlen, wpnt0, wpnt1 + + +def muscle_gain_length( + length: jax.Array, lmin: jax.Array, lmax: jax.Array +) -> jax.Array: + """Normalized muscle length-gain curve.""" + # mid-ranges (maximum is at 1.0) + a = 0.5 * (lmin + 1) + b = 0.5 * (1 + lmax) + + out0 = 0.5 * jp.square( + (length - lmin) / jp.maximum(mujoco.mjMINVAL, a - lmin) + ) + out1 = 1 - 0.5 * jp.square((1 - length) / jp.maximum(mujoco.mjMINVAL, 1 - a)) + out2 = 1 - 0.5 * jp.square((length - 1) / jp.maximum(mujoco.mjMINVAL, b - 1)) + out3 = 0.5 * jp.square( + (lmax - length) / jp.maximum(mujoco.mjMINVAL, lmax - b) + ) + + out = jp.where(length <= b, out2, out3) + out = jp.where(length <= 1, out1, out) + out = jp.where(length <= a, out0, out) + out = jp.where((lmin <= length) & (length <= lmax), out, 0.0) + + return out + + +def muscle_gain( + length: jax.Array, + vel: jax.Array, + lengthrange: jax.Array, + acc0: jax.Array, + prm: jax.Array, +) -> jax.Array: + """Muscle active force.""" + # unpack parameters + lrange = prm[:2] + force, scale, lmin, lmax, vmax, _, fvmax = prm[2:9] + + force = jp.where(force < 0, scale / jp.maximum(mujoco.mjMINVAL, acc0), force) + + # optimum length + L0 = (lengthrange[1] - lengthrange[0]) / jp.maximum( # pylint:disable=invalid-name + mujoco.mjMINVAL, lrange[1] - lrange[0] + ) + + # normalized length and velocity + L = lrange[0] + (length - lengthrange[0]) / jp.maximum(mujoco.mjMINVAL, L0) # pylint:disable=invalid-name + V = vel / jp.maximum(mujoco.mjMINVAL, L0 * vmax) # pylint:disable=invalid-name + + # length curve + FL = muscle_gain_length(L, lmin, lmax) # pylint:disable=invalid-name + + # velocity curve + y = fvmax - 1 + FV = fvmax # pylint:disable=invalid-name + FV = jp.where( # pylint:disable=invalid-name + V <= y, fvmax - jp.square(y - V) / jp.maximum(mujoco.mjMINVAL, y), FV + ) + FV = jp.where(V <= 0, jp.square(V + 1), FV) # pylint:disable=invalid-name + FV = jp.where(V <= -1, 0, FV) # pylint:disable=invalid-name + + # compute FVL and scale, make it negative + return -force * FL * FV + + +def muscle_bias( + length: jax.Array, lengthrange: jax.Array, acc0: jax.Array, prm: jax.Array +) -> jax.Array: + """Muscle passive force.""" + # unpack parameters + lrange = prm[:2] + force, scale, _, lmax, _, fpmax = prm[2:8] + + force = jp.where(force < 0, scale / jp.maximum(mujoco.mjMINVAL, acc0), force) + + # optimum length + L0 = (lengthrange[1] - lengthrange[0]) / jp.maximum( # pylint:disable=invalid-name + mujoco.mjMINVAL, lrange[1] - lrange[0] + ) + + # normalized length + L = lrange[0] + (length - lengthrange[0]) / jp.maximum(mujoco.mjMINVAL, L0) # pylint:disable=invalid-name + + # half-quadratic to (L0 + lmax) / 2, linear beyond + b = 0.5 * (1 + lmax) + + out1 = ( + -force + * fpmax + * 0.5 + * jp.square((L - 1) / jp.maximum(mujoco.mjMINVAL, b - 1)) + ) + out2 = -force * fpmax * (0.5 + (L - b) / jp.maximum(mujoco.mjMINVAL, b - 1)) + + out = jp.where(L <= b, out1, out2) + out = jp.where(L <= 1, 0.0, out) + + return out + + +def muscle_dynamics_timescale( + dctrl: jax.Array, + tau_act: jax.Array, + tau_deact: jax.Array, + smoothing_width: jax.Array, +) -> jax.Array: + """Muscle time constant with optional smoothing.""" + # hard switching + tau_hard = jp.where(dctrl > 0, tau_act, tau_deact) + + def _sigmoid(x): + # sigmoid function over 0 <= x <= 1 using quintic polynomial + # sigmoid: f(x) = 6 * x^5 - 15 * x^4 + 10 * x^3 + # solution of f(0) = f'(0) = f''(0) = 0, f(1) = 1, f'(1) = f''(1) = 0 + return jp.clip(x**3 * (3 * x * (2 * x - 5) + 10), 0, 1) + + # smooth switching + # scale by width, center around 0.5 midpoint, rescale to bounds + tau_smooth = tau_deact + (tau_act - tau_deact) * _sigmoid( + dctrl / smoothing_width + 0.5 + ) + + return jp.where(smoothing_width < mujoco.mjMINVAL, tau_hard, tau_smooth) + + +def muscle_dynamics( + ctrl: jax.Array, act: jax.Array, prm: jax.Array +) -> jax.Array: + """Muscle activation dynamics.""" + # clamp control + ctrlclamp = jp.clip(ctrl, 0, 1) + + # clamp activation + actclamp = jp.clip(act, 0, 1) + + # compute timescales as in Millard et at. (2013) + # https://doi.org/10.1115/1.4023390 + tau_act = prm[0] * (0.5 + 1.5 * actclamp) # activation timescale + tau_deact = prm[1] / (0.5 + 1.5 * actclamp) # deactivation timescale + smoothing_width = prm[2] # width of smoothing sigmoid + dctrl = ctrlclamp - act # excess excitation + + tau = muscle_dynamics_timescale(dctrl, tau_act, tau_deact, smoothing_width) + + # filter output + return dctrl / jp.maximum(mujoco.mjMINVAL, tau) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 1263f1a7..c955ba72 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -221,6 +221,211 @@ class SupportTest(parameterized.TestCase): force = force.at[3:].set(dx.contact.frame[j] @ force[3:]) np.testing.assert_allclose(result, force, rtol=1e-5, atol=2) + def test_muscle_gain_length(self): + lmin = 0.5 + lmax = 1.5 + np.testing.assert_allclose( + support.muscle_gain_length(0, lmin, lmax), + jp.zeros(1), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + support.muscle_gain_length(0.5, lmin, lmax), + jp.zeros(1), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + support.muscle_gain_length(0.6, lmin, lmax), + jp.array([0.08]), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + support.muscle_gain_length(0.75, lmin, lmax), + jp.array([0.5]), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + support.muscle_gain_length(1.0, lmin, lmax), + jp.ones(1), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + support.muscle_gain_length(1.25, lmin, lmax), + jp.array([0.5]), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + support.muscle_gain_length(1.5, lmin, lmax), + jp.zeros(1), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + support.muscle_gain_length(2.0, lmin, lmax), + jp.zeros(1), + rtol=1e-5, + atol=1e-5, + ) + + def test_muscle_gain(self): + length = jp.array([1.0]) + lengthrange = jp.array([0.0, 1.0]) + acc0 = jp.array([1.0]) + prm = jp.array([0.0, 1.0, 1.0, 200.0, 0.5, 3.0, 1.0, 0.0, 2.0, 0.0]) + + # V <= -1 + vel = jp.array([-1.5]) + np.testing.assert_allclose( + support.muscle_gain(length, vel, lengthrange, acc0, prm), + jp.array([-0.0]), + rtol=1e-5, + atol=1e-5, + ) + + # V <= 0 + vel = jp.array([-0.5]) + np.testing.assert_allclose( + support.muscle_gain(length, vel, lengthrange, acc0, prm), + jp.array([-0.25]), + rtol=1e-5, + atol=1e-5, + ) + + # V <= y + vel = jp.array([0.5]) + np.testing.assert_allclose( + support.muscle_gain(length, vel, lengthrange, acc0, prm), + jp.array([-1.75]), + rtol=1e-5, + atol=1e-5, + ) + + # V > y + vel = jp.array([1.5]) + np.testing.assert_allclose( + support.muscle_gain(length, vel, lengthrange, acc0, prm), + jp.array([-2.0]), + rtol=1e-5, + atol=1e-5, + ) + + # force < 0 + prm = prm.at[2].set(-1.0) + np.testing.assert_allclose( + support.muscle_gain(length, vel, lengthrange, acc0, prm), + jp.array([-400.0]), + rtol=1e-5, + atol=1e-5, + ) + + def test_muscle_bias(self): + lengthrange = jp.array([0.0, 1.0]) + acc0 = jp.array([1.0]) + prm = jp.array([0.0, 1.0, 1.0, 200.0, 0.5, 3.0, 1.5, 1.3, 1.2, 0.0]) + + # L <= 1 + length = jp.array([0.5]) + np.testing.assert_allclose( + support.muscle_bias(length, lengthrange, acc0, prm), + jp.array([0.0]), + rtol=1e-5, + atol=1e-5, + ) + + # L <= b + length = jp.array([1.5]) + np.testing.assert_allclose( + support.muscle_bias(length, lengthrange, acc0, prm), + jp.array([-0.1625]), + rtol=1e-5, + atol=1e-5, + ) + + # L > b + length = jp.array([2.5]) + np.testing.assert_allclose( + support.muscle_bias(length, lengthrange, acc0, prm), + jp.array([-1.3]), + rtol=1e-5, + atol=1e-5, + ) + + # force < 0 + prm = prm.at[2].set(-1.0) + np.testing.assert_allclose( + support.muscle_bias(length, lengthrange, acc0, prm), + jp.array([-260.0]), + rtol=1e-5, + atol=1e-5, + ) + + def test_smooth_muscle_dynamics(self): + # compute time constant as in Millard et al. (2013) + # https://doi.org/10.1115/1.4023390 + def _muscle_dynamics_millard(ctrl, act, prm): + ctrlclamp = jp.clip(ctrl, 0, 1) + actclamp = jp.clip(act, 0, 1) + + tau0 = prm[0] * (0.5 + 1.5 * actclamp) + tau1 = prm[1] / (0.5 + 1.5 * actclamp) + tau = jp.where(ctrlclamp > act, tau0, tau1) + + return (ctrlclamp - act) / jp.maximum(mujoco.mjMINVAL, tau) + + prm = jp.array([0.01, 0.04, 0.0]) + + # exact equality if tau_smooth = 0 + for ctrl in [-0.1, 0.0, 0.4, 0.5, 1.0, 1.0]: + for act in [-0.1, 0.0, 0.4, 0.5, 1.0, 1.1]: + actdot_old = _muscle_dynamics_millard(ctrl, act, prm) + actdot_new = support.muscle_dynamics(ctrl, act, prm) + np.testing.assert_allclose(actdot_old, actdot_new, rtol=1e-5, atol=1e-5) + + # positive tau_smooth + tau_smooth = 0.2 + prm = prm.at[2].set(tau_smooth) + act = 0.5 + eps = 1.0e-6 + + ctrl = 0.4 - eps # smaller than act by just over 0.5 * tau_smooth + np.testing.assert_allclose( + _muscle_dynamics_millard(ctrl, act, prm), + support.muscle_dynamics(ctrl, act, prm), + rtol=1e-5, + atol=1e-5, + ) + + ctrl = 0.6 + eps # larger than act by just over 0.5 * tau_smooth + np.testing.assert_allclose( + _muscle_dynamics_millard(ctrl, act, prm), + support.muscle_dynamics(ctrl, act, prm), + rtol=1e-5, + atol=1e-5, + ) + + # right in the middle should give average of time constants + tau_act = 0.2 + tau_deact = 0.3 + for dctrl in [0.0, 0.1, 0.2, 1.0, 1.1]: + lower = support.muscle_dynamics_timescale( + -dctrl, tau_act, tau_deact, tau_smooth + ) + upper = support.muscle_dynamics_timescale( + dctrl, tau_act, tau_deact, tau_smooth + ) + np.testing.assert_allclose( + 0.5 * (upper + lower), + 0.5 * (tau_act + tau_deact), + rtol=1e-5, + atol=1e-5, + ) + if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index ec17efed..a940d44e 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -234,12 +234,14 @@ class DynType(enum.IntEnum): INTEGRATOR: integrator: da/dt = u FILTER: linear filter: da/dt = (u-a) / tau FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration + MUSCLE: piece-wise linear filter with two time constants """ NONE = mujoco.mjtDyn.mjDYN_NONE INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR FILTER = mujoco.mjtDyn.mjDYN_FILTER FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT - # unsupported: MUSCLE, USER + MUSCLE = mujoco.mjtDyn.mjDYN_MUSCLE + # unsupported: USER class GainType(enum.IntEnum): @@ -248,10 +250,12 @@ class GainType(enum.IntEnum): Members: FIXED: fixed gain AFFINE: const + kp*length + kv*velocity + MUSCLE: muscle FLV curve computed by muscle_gain """ FIXED = mujoco.mjtGain.mjGAIN_FIXED AFFINE = mujoco.mjtGain.mjGAIN_AFFINE - # unsupported: MUSCLE, USER + MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE + # unsupported: USER class BiasType(enum.IntEnum): @@ -260,10 +264,12 @@ class BiasType(enum.IntEnum): Members: NONE: no bias AFFINE: const + kp*length + kv*velocity + MUSCLE: muscle passive force computed by muscle_bias """ NONE = mujoco.mjtBias.mjBIAS_NONE AFFINE = mujoco.mjtBias.mjBIAS_AFFINE - # unsupported: MUSCLE, USER + MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE + # unsupported: USER class ConstraintType(enum.IntEnum): diff --git a/mjx/mujoco/mjx/test_data/actuator/arm21.xml b/mjx/mujoco/mjx/test_data/actuator/arm21.xml new file mode 100644 index 00000000..57ac7990 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/actuator/arm21.xml @@ -0,0 +1,37 @@ + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/actuator/arm26.xml b/mjx/mujoco/mjx/test_data/actuator/arm26.xml new file mode 100644 index 00000000..0875358a --- /dev/null +++ b/mjx/mujoco/mjx/test_data/actuator/arm26.xml @@ -0,0 +1,118 @@ + + + + diff --git a/mjx/mujoco/mjx/test_data/actuator/general_dyntype.xml b/mjx/mujoco/mjx/test_data/actuator/general_dyntype.xml new file mode 100644 index 00000000..01d0a503 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/actuator/general_dyntype.xml @@ -0,0 +1,18 @@ + + + + + + + + + + + + + + + + + +