add muscle actuator utilities

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)

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()