Add muscle actuators to MJX.

PiperOrigin-RevId: 695334356
Change-Id: I6590ed6bdd4a8adc53d5c4e2641f472f7439cedc

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(