Implicitfast integration for MJX.

PiperOrigin-RevId: 664889120
Change-Id: Id3bd3916fe821ab2af79c52e53e14837ae829b2a

mjx/mujoco/mjx/__init__.py

@@ -23,6 +23,7 @@ from mujoco.mjx._src.forward import fwd_acceleration
 from mujoco.mjx._src.forward import fwd_actuation
 from mujoco.mjx._src.forward import fwd_position
 from mujoco.mjx._src.forward import fwd_velocity
+from mujoco.mjx._src.forward import implicit
 from mujoco.mjx._src.forward import rungekutta4
 from mujoco.mjx._src.forward import step
 from mujoco.mjx._src.io import get_data
@@ -32,9 +33,9 @@ from mujoco.mjx._src.io import put_data
 from mujoco.mjx._src.io import put_model
 from mujoco.mjx._src.passive import passive
 from mujoco.mjx._src.ray import ray
+from mujoco.mjx._src.sensor import sensor_acc
 from mujoco.mjx._src.sensor import sensor_pos
 from mujoco.mjx._src.sensor import sensor_vel
-from mujoco.mjx._src.sensor import sensor_acc
 from mujoco.mjx._src.smooth import camlight
 from mujoco.mjx._src.smooth import com_pos
 from mujoco.mjx._src.smooth import com_vel

mjx/mujoco/mjx/_src/forward.py

@@ -348,6 +348,46 @@ def rungekutta4(m: Model, d: Data) -> Data:
   return d
 
 
+@named_scope
+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')
+
+  qacc = d.qacc
+  if qderiv is not None:
+    # TODO(robotics-simulation): use smooth.factor_m / solve_m here:
+    qm = support.full_m(m, d) if support.is_sparse(m) else d.qM
+    qm -= m.opt.timestep * qderiv
+    qh, _ = jax.scipy.linalg.cho_factor(qm)
+    qfrc = d.qfrc_smooth + d.qfrc_constraint
+    qacc = jax.scipy.linalg.cho_solve((qh, False), qfrc)
+
+  return _advance(m, d, d.act_dot, qacc)
+
+
 @named_scope
 def forward(m: Model, d: Data) -> Data:
   """Forward dynamics."""
@@ -377,6 +417,8 @@ def step(m: Model, d: Data) -> Data:
     d = euler(m, d)
   elif m.opt.integrator == IntegratorType.RK4:
     d = rungekutta4(m, d)
+  elif m.opt.integrator == IntegratorType.IMPLICITFAST:
+    d = implicit(m, d)
   else:
     raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.')
 

mjx/mujoco/mjx/_src/forward_test.py

@@ -68,6 +68,12 @@ class ForwardTest(absltest.TestCase):
     _assert_attr_eq(d, dx, 'qpos')
     _assert_attr_eq(d, dx, 'time')
 
+    # implicitfast
+    m.opt.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
+    dx = jax.jit(mjx.implicit)(mx, mjx.put_data(m, d))
+    mujoco.mj_implicit(m, d)
+    _assert_attr_eq(d, dx, 'qpos')
+
   def test_step(self):
     m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)

mjx/mujoco/mjx/_src/io.py

@@ -47,13 +47,18 @@ def _make_option(o: mujoco.MjOption) -> types.Option:
     if o.enableflags & 2**i:
       raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
 
+  has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
+  implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST
+  if implicitfast and has_fluid_params:
+    raise NotImplementedError('implicitfast not implemented for fluid drag.')
+
   fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()}
   fields['integrator'] = types.IntegratorType(o.integrator)
   fields['cone'] = types.ConeType(o.cone)
   fields['jacobian'] = types.JacobianType(o.jacobian)
   fields['solver'] = types.SolverType(o.solver)
   fields['disableflags'] = types.DisableBit(o.disableflags)
-  fields['has_fluid_params'] = o.density > 0 or o.viscosity > 0 or o.wind.any()
+  fields['has_fluid_params'] = has_fluid_params
 
   return types.Option(**fields)
 

mjx/mujoco/mjx/_src/io_test.py

@@ -208,6 +208,14 @@ class ModelIOTest(parameterized.TestCase):
           </worldbody>
         </mujoco>"""))
 
+  def test_implicitfast_fluid_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <option viscosity="3.0" integrator="implicitfast"/>
+          <worldbody/>
+        </mujoco>"""))
+
 
 class DataIOTest(parameterized.TestCase):
   """IO tests for mjx.Data."""

mjx/mujoco/mjx/_src/types.py

@@ -90,10 +90,12 @@ class IntegratorType(enum.IntEnum):
   Members:
     EULER: semi-implicit Euler
     RK4: 4th-order Runge Kutta
+    IMPLICITFAST: implicit in velocity, no rne derivative
   """
   EULER = mujoco.mjtIntegrator.mjINT_EULER
   RK4 = mujoco.mjtIntegrator.mjINT_RK4
-  # unsupported: IMPLICIT, IMPLICITFAST
+  IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
+  # unsupported: IMPLICIT
 
 
 class GeomType(enum.IntEnum):

mjx/mujoco/mjx/test_data/constraints.xml

@@ -83,7 +83,7 @@
   </equality>
 
   <actuator>
-    <position ctrlrange="-20 20" gear="500" joint="joint1" name="act1"/>
+    <position ctrlrange="-20 20" gear="500" joint="joint1" kv="0.5" name="act1"/>
     <motor gear="50000" joint="joint3" name="act2"/>
     <motor gear="75000" joint="joint4" name="act3"/>
   </actuator>