diff --git a/mujoco_warp/_src/util_misc.py b/mujoco_warp/_src/util_misc.py index 87efccfd..0678b6f5 100644 --- a/mujoco_warp/_src/util_misc.py +++ b/mujoco_warp/_src/util_misc.py @@ -22,6 +22,7 @@ from . import math from .types import MJ_MINVAL from .types import WrapType +from .types import vec10 @wp.func @@ -426,3 +427,155 @@ def wrap( wpnt1 = mat @ res1 + pos return wlen, wpnt0, wpnt1 + + +@wp.func +def muscle_gain_length(length: float, lmin: float, lmax: float) -> float: + """Normalized muscle length-gain curve.""" + + if (lmin > length) or (length > lmax): + return 0.0 + + # mid-ranges (maximum is at 1.0) + a = 0.5 * (lmin + 1.0) + b = 0.5 * (1.0 + lmax) + + if length <= a: + x = (length - lmin) / wp.max(MJ_MINVAL, a - lmin) + return 0.5 * x * x + elif length <= 1.0: + x = (1.0 - length) / wp.max(MJ_MINVAL, 1.0 - a) + return 1.0 - 0.5 * x * x + elif length <= b: + x = (length - 1.0) / wp.max(MJ_MINVAL, b - 1.0) + return 1.0 - 0.5 * x * x + else: + x = (lmax - length) / wp.max(MJ_MINVAL, lmax - b) + return 0.5 * x * x + + +@wp.func +def muscle_gain(len: float, vel: float, lengthrange: wp.vec2, acc0: float, prm: vec10) -> float: + """Muscle active force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax).""" + + # unpack parameters + range_ = wp.vec2(prm[0], prm[1]) + force = prm[2] + scale = prm[3] + lmin = prm[4] + lmax = prm[5] + vmax = prm[6] + fvmax = prm[8] + + # scale force if negative + if force < 0.0: + force = scale / wp.max(MJ_MINVAL, acc0) + + # optimum length + L0 = (lengthrange[1] - lengthrange[0]) / wp.max(MJ_MINVAL, range_[1] - range_[0]) + + # normalized length and velocity + L = range_[0] + (len - lengthrange[0]) / wp.max(MJ_MINVAL, L0) + V = vel / wp.max(MJ_MINVAL, L0 * vmax) + + # length curve + FL = muscle_gain_length(L, lmin, lmax) + + # velocity curve + y = fvmax - 1.0 + if V <= -1.0: + FV = 0.0 + elif V <= 0.0: + FV = (V + 1.0) * (V + 1.0) + elif V <= y: + FV = fvmax - (y - V) * (y - V) / wp.max(MJ_MINVAL, y) + else: + FV = fvmax + + # compute FVL and scale, make it negative + return -force * FL * FV + + +@wp.func +def muscle_bias(len: float, lengthrange: wp.vec2, acc0: float, prm: vec10) -> float: + """Muscle passive force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax).""" + + # unpack parameters + range_ = wp.vec2(prm[0], prm[1]) + force = prm[2] + scale = prm[3] + lmax = prm[5] + fpmax = prm[7] + + # scale force if negative + if force < 0.0: + force = scale / wp.max(MJ_MINVAL, acc0) + + # optimum length + L0 = (lengthrange[1] - lengthrange[0]) / wp.max(MJ_MINVAL, range_[1] - range_[0]) + + # normalized length + L = range_[0] + (len - lengthrange[0]) / wp.max(MJ_MINVAL, L0) + + # half-quadratic to (L0 + lmax) / 2, linear beyond + b = 0.5 * (1.0 + lmax) + if L <= 1.0: + return 0.0 + elif L <= b: + x = (L - 1.0) / wp.max(MJ_MINVAL, b - 1.0) + return -force * fpmax * 0.5 * x * x + else: + x = (L - b) / wp.max(MJ_MINVAL, b - 1.0) + return -force * fpmax * (0.5 + x) + + +@wp.func +def _sigmoid(x: float) -> float: + """Sigmoid function over 0 <= x <= 1 using quintic polynomial.""" + + if x <= 0.0: + return 0.0 + + if x >= 1.0: + return 1.0 + + # 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 x * x * x * (3.0 * x * (2.0 * x - 5.0) + 10.0) + + +@wp.func +def muscle_dynamics_timescale(dctrl: float, tau_act: float, tau_deact: float, smooth_width: float) -> float: + """Muscle time constant with optional smoothing.""" + + # hard switching + if smooth_width < MJ_MINVAL: + if dctrl > 0.0: + return tau_act + else: + return tau_deact + else: # smooth switching + # scale by width, center around 0.5 midpoint, rescale to bounds + return tau_deact + (tau_act - tau_deact) * _sigmoid(dctrl / smooth_width + 0.5) + + +@wp.func +def muscle_dynamics(control: float, activation: float, prm: vec10) -> float: + """Muscle activation dynamics, prm = (tau_act, tau_deact, smooth_width).""" + + # clamp control + ctrlclamp = wp.clamp(control, 0.0, 1.0) + + # clamp activation + actclamp = wp.clamp(activation, 0.0, 1.0) + + # compute timescales as in Millard et al. (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 + smooth_width = prm[2] # width of smoothing sigmoid + dctrl = ctrlclamp - activation # excess excitation + + tau = muscle_dynamics_timescale(dctrl, tau_act, tau_deact, smooth_width) + + # filter output + return dctrl / wp.max(MJ_MINVAL, tau) diff --git a/mujoco_warp/_src/util_misc_test.py b/mujoco_warp/_src/util_misc_test.py index 00930414..6bacb687 100644 --- a/mujoco_warp/_src/util_misc_test.py +++ b/mujoco_warp/_src/util_misc_test.py @@ -25,6 +25,7 @@ from . import util_misc from .types import MJ_MINVAL from .types import WrapType +from .types import vec10 def _assert_eq(a, b, name): @@ -225,6 +226,73 @@ def wrap( return length.numpy()[0], wpnt0.numpy()[0], wpnt1.numpy()[0] +def _muscle_dynamics_millard(ctrl, act, prm): + """Compute time constant as in Millard et al. (2013) https://doi.org/10.1115/1.4023390.""" + + # clamp control + ctrlclamp = np.clip(ctrl, 0.0, 1.0) + + # clamp activation + actclamp = np.clip(act, 0.0, 1.0) + + if ctrlclamp > act: + tau = prm[0] * (0.5 + 1.5 * actclamp) + else: + tau = prm[1] / (0.5 + 1.5 * actclamp) + + # filter output + return (ctrlclamp - act) / np.maximum(MJ_MINVAL, tau) + + +def _muscle_dynamics(ctrl, act, prm): + @wp.kernel + def muscle_dynamics(control: float, activation: float, prm: vec10, dynamics_out: wp.array(dtype=float)): + dynamics_out[0] = util_misc.muscle_dynamics(control, activation, prm) + + output = wp.empty(1, dtype=float) + wp.launch( + muscle_dynamics, + dim=(1,), + inputs=[ + ctrl, + act, + vec10(prm[0], prm[1], prm[2], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + ], + outputs=[output], + ) + + return output.numpy()[0] + + +def _muscle_gain_length(length, lmin, lmax): + @wp.kernel + def muscle_gain_length(length: float, lmin: float, lmax: float, gain_length_out: wp.array(dtype=float)): + gain_length_out[0] = util_misc.muscle_gain_length(length, lmin, lmax) + + output = wp.empty(1, dtype=float) + wp.launch(muscle_gain_length, dim=(1,), inputs=[length, lmin, lmax], outputs=[output]) + + return output.numpy()[0] + + +def _muscle_dynamics_timescale(dctrl, tau_act, tau_deact, smooth_width): + @wp.kernel + def muscle_gain_length( + dctrl: float, tau_act: float, tau_deact: float, smooth_width: float, dynamics_timescale_out: wp.array(dtype=float) + ): + dynamics_timescale_out[0] = util_misc.muscle_dynamics_timescale(dctrl, tau_act, tau_deact, smooth_width) + + output = wp.empty(1, dtype=float) + wp.launch( + muscle_gain_length, + dim=(1,), + inputs=[dctrl, tau_act, tau_deact, smooth_width], + outputs=[output], + ) + + return output.numpy()[0] + + class UtilMiscTest(parameterized.TestCase): def test_is_intersect(self): self.assertFalse( @@ -442,6 +510,63 @@ def test_wrap(self, wraptype): _assert_eq(wpnt0, np.array([0.1, 0.0, 0.0]), "wpnt0") _assert_eq(wpnt1, np.array([0.1, 0.0, 0.0]), "wpnt1") + @parameterized.product(ctrl=[-0.1, 0.0, 0.4, 0.5, 1.0, 1.1], act=[-0.1, 0.0, 0.4, 0.5, 1.0, 1.1]) + def test_muscle_dynamics_tausmooth0(self, ctrl, act): + # exact equality if tau_smooth = 0 + prm = np.array([0.01, 0.04, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + + actdot_old = _muscle_dynamics_millard(ctrl, act, prm) + actdot_new = _muscle_dynamics(ctrl, act, prm) + + _assert_eq(actdot_new, actdot_old, "actdot") + + def test_muscle_dynamics_tausmooth_positive(self): + # positive tau_smooth + prm = np.array([0.01, 0.04, 0.2, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + act = 0.5 + eps = 1.0e-6 + + ctrl = 0.4 - eps # smaller than act by just over 0.5 * tau_smooth + _assert_eq( + _muscle_dynamics(ctrl, act, prm), + _muscle_dynamics_millard(ctrl, act, prm), + "actdot", + ) + + ctrl = 0.6 + eps # larger than act by just over 0.5 * tau_smooth + _assert_eq( + _muscle_dynamics(ctrl, act, prm), + _muscle_dynamics_millard(ctrl, act, prm), + "actdot", + ) + + @parameterized.parameters(0.0, 0.1, 0.2, 1.0, 1.1) + def test_muscle_dynamics_timescale(self, dctrl): + # right in the middle should give average of time constants + tau_smooth = 0.2 + tau_act = 0.2 + tau_deact = 0.3 + + lower = _muscle_dynamics_timescale(-dctrl, tau_act, tau_deact, tau_smooth) + upper = _muscle_dynamics_timescale(dctrl, tau_act, tau_deact, tau_smooth) + + _assert_eq(0.5 * (lower + upper), 0.5 * (tau_act + tau_deact), "muscle_dynamics_timescale") + + @parameterized.parameters( + (0.0, 0.0), + (0.5, 0.0), + (0.75, 0.5), + (1.0, 1.0), + (1.25, 0.5), + (1.5, 0.0), + (2.0, 0.0), + ) + def test_muscle_gain_length(self, input, output): + _assert_eq(_muscle_gain_length(input, 0.5, 1.5), output, "length-gain") + + # TODO(team): test util_misc.muscle_gain + # TODO(team): test util_misc.muscle_bias + if __name__ == "__main__": wp.init()