Add inverse dynamics to MJX.

PiperOrigin-RevId: 745998273
Change-Id: I203af89332ace5d7a60fff3ec6ff4cc88c02340d

mjx/mujoco/mjx/_src/forward.py

@@ -22,6 +22,7 @@ from jax import numpy as jp
 import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import constraint
+from mujoco.mjx._src import derivative
 from mujoco.mjx._src import math
 from mujoco.mjx._src import passive
 from mujoco.mjx._src import scan
@@ -392,29 +393,7 @@ def rungekutta4(m: Model, d: Data) -> Data:
 def implicit(m: Model, d: Data) -> Data:
   """Integrates fully implicit in velocity."""
 
-  qderiv = None
-
-  # qDeriv += d qfrc_actuator / d qvel
-  if not m.opt.disableflags & DisableBit.ACTUATION:
-    affine_bias = m.actuator_biastype == BiasType.AFFINE
-    bias_vel = m.actuator_biasprm[:, 2] * affine_bias
-    affine_gain = m.actuator_gaintype == GainType.AFFINE
-    gain_vel = m.actuator_gainprm[:, 2] * affine_gain
-    ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act)
-    vel = bias_vel + gain_vel * ctrl
-    qderiv = d.actuator_moment.T @ jp.diag(vel) @ d.actuator_moment
-
-  # qDeriv += d qfrc_passive / d qvel
-  if not m.opt.disableflags & DisableBit.PASSIVE:
-    if qderiv is None:
-      qderiv = -jp.diag(m.dof_damping)
-    else:
-      qderiv -= jp.diag(m.dof_damping)
-    if m.ntendon:
-      qderiv -= d.ten_J.T @ jp.diag(m.tendon_damping) @ d.ten_J
-    # TODO(robotics-simulation): fluid drag model
-    if m.opt.has_fluid_params:
-      raise NotImplementedError('fluid drag not supported for implicitfast')
+  qderiv = derivative.deriv_smooth_vel(m, d)
 
   qacc = d.qacc
   if qderiv is not None: