Add inverse dynamics to MJX.

PiperOrigin-RevId: 745998273
Change-Id: I203af89332ace5d7a60fff3ec6ff4cc88c02340d

Author: thowell

doc/changelog.rst

@@ -1,6 +1,14 @@
 =========
 Changelog
 =========
+
+Upcoming version (not yet release)
+----------------------------------
+
+MJX
+^^^
+- Added inverse dynamics.
+
 Version 3.3.1 (Apr 9, 2025)
 ----------------------------
 

doc/mjx.rst

@@ -235,6 +235,8 @@ The following features are **fully supported** in MJX:
      - 1, 3, 4, 6 (1 is not supported with ``ELLIPTIC``)
    * - :ref:`Solver <mjtSolver>`
      - ``CG``, ``NEWTON``
+   * - Dynamics
+     - :ref:`Inverse <mj_inverse>`
    * - Fluid Model
      - :ref:`flInertia`
    * - :ref:`Tendons <tendon>`
@@ -262,8 +264,6 @@ The following features are **in development** and coming soon:
        (``BOX``, ``MESH``, ``HFIELD``) and ``ELLIPSOID``.
    * - :ref:`Integrator <mjtIntegrator>`
      - ``IMPLICIT``
-   * - Dynamics
-     - :ref:`Inverse <mj_inverse>`
    * - Fluid Model
      - :ref:`flEllipsoid`
    * - :ref:`Sensors <mjtSensor>`

mjx/mujoco/mjx/__init__.py

@@ -17,6 +17,7 @@
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
 from mujoco.mjx._src.constraint import make_constraint
+from mujoco.mjx._src.derivative import deriv_smooth_vel
 from mujoco.mjx._src.forward import euler
 from mujoco.mjx._src.forward import forward
 from mujoco.mjx._src.forward import fwd_acceleration
@@ -26,6 +27,7 @@
 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.inverse import inverse
 from mujoco.mjx._src.io import get_data
 from mujoco.mjx._src.io import get_data_into
 from mujoco.mjx._src.io import make_data

mjx/mujoco/mjx/_src/derivative.py

@@ -0,0 +1,59 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Derivative functions."""
+
+from typing import Optional
+
+import jax
+from jax import numpy as jp
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import BiasType
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DisableBit
+from mujoco.mjx._src.types import DynType
+from mujoco.mjx._src.types import GainType
+from mujoco.mjx._src.types import Model
+
+
+def deriv_smooth_vel(m: Model, d: Data) -> Optional[jax.Array]:
+  """Analytical derivative of smooth forces w.r.t velocities."""
+
+  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 @ jax.vmap(jp.multiply)(d.actuator_moment, vel)
+
+  # 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')
+
+  # TODO(team): rne derivative
+
+  return qderiv

mjx/mujoco/mjx/_src/forward.py

@@ -22,6 +22,7 @@
 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 f(carry, x):
 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:

mjx/mujoco/mjx/_src/inverse.py

@@ -0,0 +1,106 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Inverse dynamics functions."""
+
+from jax import numpy as jp
+from mujoco.mjx._src import derivative
+from mujoco.mjx._src import forward
+from mujoco.mjx._src import sensor
+from mujoco.mjx._src import smooth
+from mujoco.mjx._src import solver
+from mujoco.mjx._src import support
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DisableBit
+from mujoco.mjx._src.types import EnableBit
+from mujoco.mjx._src.types import IntegratorType
+from mujoco.mjx._src.types import Model
+
+
+def discrete_acc(m: Model, d: Data) -> Data:
+  """Convert discrete-time qacc to continuous-time qacc."""
+
+  if m.opt.integrator == IntegratorType.RK4:
+    raise RuntimeError(
+        'discrete inverse dynamics is not supported by RK4 integrator'
+    )
+  elif m.opt.integrator == IntegratorType.EULER:
+    dsbl_eulerdamp = m.opt.disableflags & DisableBit.EULERDAMP
+    no_dof_damping = (m.dof_damping == 0).all()
+    if dsbl_eulerdamp or no_dof_damping:
+      return d
+
+    # set qfrc = (M + h*diag(B)) * qacc
+    qfrc = support.mul_m(m, d, d.qacc)
+    qfrc += m.opt.timestep * m.dof_damping * d.qacc
+  elif m.opt.integrator == IntegratorType.IMPLICITFAST:
+    qm = support.full_m(m, d)
+
+    # compute analytical derivative qDeriv; skip rne derivative
+    qderiv = derivative.deriv_smooth_vel(m, d)
+    if qderiv is not None:
+      # M = M - dt*qDeriv
+      qm -= m.opt.timestep * qderiv
+
+    # set qfrc = (M - dt*qDeriv) * qacc
+    qfrc = qm @ d.qacc
+  else:
+    raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.')
+
+  # solve for qacc: qfrc = M * qacc
+  qacc = smooth.solve_m(m, d, qfrc)
+
+  return d.replace(qacc=qacc)
+
+
+def inv_constraint(m: Model, d: Data) -> Data:
+  """Inverse constraint solver."""
+
+  # no constraints
+  if d.efc_J.size == 0:
+    return d.replace(qfrc_constraint=jp.zeros(m.nv))
+
+  # update
+  ctx = solver.Context.create(m, d, grad=False)
+
+  return d.replace(
+      qfrc_constraint=ctx.qfrc_constraint,
+      efc_force=ctx.efc_force,
+  )
+
+
+def inverse(m: Model, d: Data) -> Data:
+  """Inverse dynamics."""
+  d = forward.fwd_position(m, d)
+  d = sensor.sensor_pos(m, d)
+  d = forward.fwd_velocity(m, d)
+  d = sensor.sensor_vel(m, d)
+
+  qacc = d.qacc
+  if m.opt.enableflags & EnableBit.INVDISCRETE:
+    d = discrete_acc(m, d)
+
+  d = inv_constraint(m, d)
+  d = smooth.rne(m, d, flg_acc=True)
+  d = sensor.sensor_acc(m, d)
+
+  qfrc_inverse = (
+      d.qfrc_bias + m.dof_armature * d.qacc - d.qfrc_passive - d.qfrc_constraint
+  )
+
+  if m.opt.enableflags & EnableBit.INVDISCRETE:
+    return d.replace(qfrc_inverse=qfrc_inverse, qacc=qacc)
+  else:
+    return d.replace(qfrc_inverse=qfrc_inverse)

mjx/mujoco/mjx/_src/inverse_test.py

@@ -0,0 +1,131 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for inverse dynamics functions."""
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax import numpy as jp
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import support
+from mujoco.mjx._src import test_util
+import numpy as np
+
+# tolerance for difference between MuJoCo and MJX calculations - mostly
+# due to float precision
+_TOLERANCE = 1e-5
+
+
+def _assert_eq(a, b, name, tol=_TOLERANCE):
+  tol = tol * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class InverseTest(parameterized.TestCase):
+
+  @parameterized.parameters(
+      (mujoco.mjtIntegrator.mjINT_EULER, False, False),
+      (mujoco.mjtIntegrator.mjINT_EULER, False, True),
+      (mujoco.mjtIntegrator.mjINT_EULER, True, False),
+      (mujoco.mjtIntegrator.mjINT_EULER, True, True),
+      (mujoco.mjtIntegrator.mjINT_IMPLICITFAST, False, False),
+      (mujoco.mjtIntegrator.mjINT_IMPLICITFAST, True, False),
+  )
+  def test_forward_inverse_match(self, integrator, invdiscrete, eulerdamp):
+    m = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <option timestep=".005" gravity="-1 -1 -10"/>
+        <worldbody>
+          <geom type="plane" size="10 10 .001"/>
+          <body pos="0 0 1">
+            <geom type="sphere" size=".1" pos=".1 .2 .3"/>
+            <joint name="jnt1" type="hinge" axis="0 1 0" stiffness=".25" damping=".125"/>
+            <body pos="0 0 1">
+              <geom type="sphere" size=".1" pos=".1 .2 .3"/>
+              <joint name="jnt2" type="hinge" axis="0 1 0" stiffness=".6" damping=".3"/>
+            </body>
+          </body>
+        </worldbody>
+        <actuator>
+          <motor joint="jnt1"/>
+        </actuator>
+        <equality>
+          <joint joint1="jnt1" joint2="jnt2"/>
+        </equality>
+      </mujoco>
+    """)
+    m.opt.integrator = integrator
+    if invdiscrete:
+      m.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_INVDISCRETE
+    if not eulerdamp:
+      m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_EULERDAMP
+
+    d = mujoco.MjData(m)
+    d.qvel = np.random.uniform(low=-0.01, high=0.01, size=d.qvel.shape)
+    d.ctrl = np.random.uniform(low=-0.01, high=0.01, size=d.ctrl.shape)
+    d.qfrc_applied = np.random.uniform(
+        low=-0.01, high=0.01, size=d.qfrc_applied.shape
+    )
+    d.xfrc_applied = np.random.uniform(
+        low=-0.01, high=0.01, size=d.xfrc_applied.shape
+    )
+    mujoco.mj_step(m, d, 100)
+
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    dx_next = mjx.step(mx, dx)
+    qacc_fd = (dx_next.qvel - dx.qvel) / mx.opt.timestep
+
+    dx = mjx.forward(mx, dx)
+
+    if invdiscrete:
+      dx = dx.replace(qacc=qacc_fd)
+
+    dxinv = mjx.inverse(mx, dx)
+
+    fwdinv0 = jp.linalg.norm(
+        dxinv.qfrc_constraint - dx.qfrc_constraint, ord=np.inf
+    )
+    fwdinv1 = jp.linalg.norm(
+        dxinv.qfrc_inverse
+        - (
+            dx.qfrc_applied + dx.qfrc_actuator + support.xfrc_accumulate(mx, dx)
+        ),
+        ord=np.inf,
+    )
+
+    self.assertLess(fwdinv0, 1.0e-3)
+    self.assertLess(fwdinv1, 1.0e-3)
+    _assert_eq(dxinv.qacc, dx.qacc, 'qacc')
+
+  def test_tendon_force_clamp(self):
+    m = test_util.load_test_file('actuator/tendon_force_clamp.xml')
+    d = mujoco.MjData(m)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    dx = dx.replace(ctrl=jp.array([1.0, 1.0, 1.0, -1.0, 1.0, -20.0, 5.0, -5.0]))
+    dx = mjx.forward(mx, dx)
+
+    _assert_eq(
+        dx.actuator_force,
+        jp.array([1.0, 1.0, 1.0, -1.0, 1.0, -10.0, 5.0, -5.0]),
+        'actuator_force',
+    )
+
+
+if __name__ == '__main__':
+  absltest.main()