Add mjd_quatIntegrate, expose mjd_subQuat.

PiperOrigin-RevId: 546252926
Change-Id: I7d6b2bd6536e7e8afc368f3c46037c20f6a1cb5e

doc/APIreference/APIfunctions.rst

@@ -29,7 +29,7 @@ API function can be classified as:
 - :ref:`Poses transformations<Poses>`.
 - :ref:`Matrix decompositions and solvers<Decompositions>`.
 - :ref:`Miscellaneous<Miscellaneous>` functions.
-- :ref:`Dynamics derivatives<Derivatives-api>`.
+- :ref:`Derivatives<Derivatives-api>`.
 - :ref:`Plugin<Plugins-api>` related functions.
 - :ref:`Macros<Macros>`.
 

doc/APIreference/functions.rst

@@ -3185,6 +3185,10 @@ Sigmoid function over 0<=x<=1 using quintic polynomial.
 Derivatives
 ^^^^^^^^^^^
 
+The functions below provide useful derivatives of various functions, both analytic and
+finite-differenced. The latter have names with the suffix ``FD``. Note that unlike much of the API,
+outputs of derivative functions are the trailing rather than leading arguments.
+
 .. _mjd_transitionFD:
 
 mjd_transitionFD
@@ -3251,6 +3255,42 @@ using finite-differencing. These matrices and their dimensions are:
 - ``flg_actuation`` denotes whether to subtract actuation forces (``qfrc_actuator``) from the output of the inverse
   dynamics. If this flag is positive, actuator forces are not considered as external.
 
+.. _mjd_subQuat:
+
+mjd_subQuat
+~~~~~~~~~~~
+
+.. mujoco-include:: mjd_subQuat
+
+Derivatives of :ref:`mju_subQuat` (quaternion difference).
+
+.. _mjd_quatIntegrate:
+
+mjd_quatIntegrate
+~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjd_quatIntegrate
+
+Derivatives of :ref:`mju_quatIntegrate`.
+
+:math:`{\tt \small mju\_quatIntegrate}(q, v, h)` performs the in-place rotation :math:`q \leftarrow q + v h`,
+where :math:`q \in \mathbf{S}^3` is a unit quaternion, :math:`v \in \mathbf{R}^3` is a 3D angular velocity and
+:math:`h \in \mathbf{R^+}` is a timestep. This is equivalent to :math:`{\tt \small mju\_quatIntegrate}(q, s, 1.0)`,
+where :math:`s` is the scaled velocity :math:`s = h v`.
+
+:math:`{\tt \small mjd\_quatIntegrate}(v, h, D_q, D_v, D_h)` computes the Jacobians of the output :math:`q` with respect
+to the inputs. Below, :math:`\bar q` denotes the pre-modified quaternion:
+
+.. math::
+   \begin{aligned}
+      D_q &= \partial q / \partial \bar q \\
+      D_v &= \partial q / \partial v \\
+      D_h &= \partial q / \partial h
+   \end{aligned}
+
+Note that derivatives depend only on :math:`h` and :math:`v` (in fact, on :math:`s = h v`).
+All outputs are optional.
+
 .. _Plugins-api:
 
 Plugins

doc/APIreference/functions_override.rst

@@ -435,6 +435,10 @@ Symmetrize square matrix :math:`R = \frac{1}{2}(M + M^T)`.
 
 .. _Derivatives-api:
 
+The functions below provide useful derivatives of various functions, both analytic and
+finite-differenced. The latter have names with the suffix ``FD``. Note that unlike much of the API,
+outputs of derivative functions are the trailing rather than leading arguments.
+
 .. _mjd_transitionFD:
 
 Finite-differenced discrete-time transition matrices.
@@ -490,3 +494,29 @@ using finite-differencing. These matrices and their dimensions are:
 - ``eps`` is the (forward) finite-differencing epsilon.
 - ``flg_actuation`` denotes whether to subtract actuation forces (``qfrc_actuator``) from the output of the inverse
   dynamics. If this flag is positive, actuator forces are not considered as external.
+
+.. _mjd_subQuat:
+
+Derivatives of :ref:`mju_subQuat` (quaternion difference).
+
+.. _mjd_quatIntegrate:
+
+Derivatives of :ref:`mju_quatIntegrate`.
+
+:math:`{\tt \small mju\_quatIntegrate}(q, v, h)` performs the in-place rotation :math:`q \leftarrow q + v h`,
+where :math:`q \in \mathbf{S}^3` is a unit quaternion, :math:`v \in \mathbf{R}^3` is a 3D angular velocity and
+:math:`h \in \mathbf{R^+}` is a timestep. This is equivalent to :math:`{\tt \small mju\_quatIntegrate}(q, s, 1.0)`,
+where :math:`s` is the scaled velocity :math:`s = h v`.
+
+:math:`{\tt \small mjd\_quatIntegrate}(v, h, D_q, D_v, D_h)` computes the Jacobians of the output :math:`q` with respect
+to the inputs. Below, :math:`\bar q` denotes the pre-modified quaternion:
+
+.. math::
+   \begin{aligned}
+      D_q &= \partial q / \partial \bar q \\
+      D_v &= \partial q / \partial v \\
+      D_h &= \partial q / \partial h
+   \end{aligned}
+
+Note that derivatives depend only on :math:`h` and :math:`v` (in fact, on :math:`s = h v`).
+All outputs are optional.

doc/changelog.rst

@@ -8,9 +8,12 @@ Upcoming version (not yet released)
 General
 ^^^^^^^
 
+- Added primitive collider for sphere-cylinder contacts, previously this pair used the generic convex-convex collider.
+- Added analytic derivatives for quaternion :ref:`subtraction<mjd_subQuat>` and :ref:`integration<mjd_quatIntegrate>`
+  (rotation with an angular velocity). Derivatives are in the 3D tangent space.
 - Added :ref:`mjv_connector` which has identical functionality to :ref:`mjv_makeConnector`, but with more convenient
   "from-to" argument parametrization. :ref:`mjv_makeConnector` is now deprecated.
-- Added primitive collider for sphere-cylinder contacts, previously this pair used the generic convex-convex collider.
+
 
 Version 2.3.6 (June 20, 2023)
 -----------------------------

doc/includes/references.h

@@ -2494,6 +2494,9 @@ void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuatio
                    mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa,
                    mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
                    mjtNum *DmDq);
+void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db[9]);
+void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
+                       mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]);
 void mjp_defaultPlugin(mjpPlugin* plugin);
 int mjp_registerPlugin(const mjpPlugin* plugin);
 int mjp_pluginCount();

include/mujoco/mujoco.h

@@ -1246,6 +1246,13 @@ MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_ac
                          mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
                          mjtNum *DmDq);
 
+// Derivatives of mju_subQuat.
+MJAPI void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db[9]);
+
+// Derivatives of mju_quatIntegrate.
+MJAPI void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
+                             mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]);
+
 //---------------------- Plugins -------------------------------------------------------------------
 
 // Set default plugin definition.

introspect/functions.py

@@ -7973,6 +7973,82 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Finite differenced Jacobians of (force, sensors) = mj_inverse(state, acceleration)   All outputs are optional. Output dimensions (transposed w.r.t Control Theory convention):     DfDq: (nv x nv)     DfDv: (nv x nv)     DfDa: (nv x nv)     DsDq: (nv x nsensordata)     DsDv: (nv x nsensordata)     DsDa: (nv x nsensordata)     DmDq: (nv x nM)   single-letter shortcuts:     inputs: q=qpos, v=qvel, a=qacc     outputs: f=qfrc_inverse, s=sensordata, m=qM   notes:     optionally computes mass matrix Jacobian DmDq     flg_actuation specifies whether to subtract qfrc_actuator from qfrc_inverse',  # pylint: disable=line-too-long
      )),
+    ('mjd_subQuat',
+     FunctionDecl(
+         name='mjd_subQuat',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='qa',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                     extents=(4,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='qb',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                     extents=(4,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='Da',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(9,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='Db',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(9,),
+                 ),
+             ),
+         ),
+         doc='Derivatives of mju_subQuat.',
+     )),
+    ('mjd_quatIntegrate',
+     FunctionDecl(
+         name='mjd_quatIntegrate',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='vel',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                     extents=(3,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='scale',
+                 type=ValueType(name='mjtNum'),
+             ),
+             FunctionParameterDecl(
+                 name='Dquat',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(9,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='Dvel',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(9,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='Dscale',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(3,),
+                 ),
+             ),
+         ),
+         doc='Derivatives of mju_quatIntegrate.',
+     )),
     ('mjp_defaultPlugin',
      FunctionDecl(
          name='mjp_defaultPlugin',

python/mujoco/bindings_test.py

@@ -1113,6 +1113,21 @@ Euler integrator, semi-implicit in velocity.
     self.assertGreater(np.linalg.norm(ds_dv), eps)
     self.assertGreater(np.linalg.norm(ds_da), eps)
 
+  def test_mjd_sub_quat(self):
+    quat1 = np.array((0.2, 0.3, 0.3, 0.4))
+    quat2 = np.array((0.2, 0.3, 0.3, 0.4))
+    d1 = np.empty(9, np.float64)
+    d2 = np.empty(9, np.float64)
+    mujoco.mjd_subQuat(quat1, quat2, d1, d2)
+
+  def test_mjd_quat_intergrate(self):
+    scale = 0.1
+    vel = np.array((0.2, 0.3, 0.3))
+    d_quat = np.empty(9, np.float64)
+    d_vel = np.empty(9, np.float64)
+    d_h = np.empty(3, np.float64)
+    mujoco.mjd_quatIntegrate(vel, scale, d_quat, d_vel, d_h)
+
   def test_banded(self):
     n_total = 4
     n_band = 1

python/mujoco/functions.cc

@@ -1239,6 +1239,11 @@ PYBIND11_MODULE(_functions, pymodule) {
         return InterceptMjErrors(::mju_insertionSortInt)(
             res.data(), res.size());
       });
+  Def<traits::mju_Halton>(pymodule);
+  // Skipped: mju_strncpy (doesn't make sense in Python)
+  Def<traits::mju_sigmoid>(pymodule);
+
+  // Derivatives
   Def<traits::mjd_transitionFD>(
       pymodule,
       [](const raw::MjModel* m, raw::MjData* d,
@@ -1319,9 +1324,8 @@ PYBIND11_MODULE(_functions, pymodule) {
             DsDa.has_value() ? DsDa->data() : nullptr,
             DmDq.has_value() ? DmDq->data() : nullptr);
       });
-  Def<traits::mju_Halton>(pymodule);
+  Def<traits::mjd_subQuat>(pymodule);
-  // Skipped: mju_strncpy (doesn't make sense in Python)
+  Def<traits::mjd_quatIntegrate>(pymodule);
-  Def<traits::mju_sigmoid>(pymodule);
 }  // PYBIND11_MODULE NOLINT(readability/fn_size)
 }  // namespace
 }  // namespace mujoco::python

src/engine/engine_derivative.c

@@ -267,6 +267,64 @@ void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db
 
 
 
+// derivative of mju_quatIntegrate w.r.t scaled velocity
+//  reference: https://arxiv.org/abs/1711.02508, Eq. 183
+void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
+                       mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]) {
+  // scaled velocity
+  mjtNum s[3] = {scale*vel[0], scale*vel[1], scale*vel[2]};
+
+  // 3 basis matrices
+  mjtNum eye[9] = {
+    1, 0, 0,
+    0, 1, 0,
+    0, 0, 1
+  };
+  mjtNum cross[9] = {
+    0,     s[2], -s[1],
+   -s[2],  0,     s[0],
+    s[1], -s[0],  0
+  };
+  mjtNum outer[9] = {
+    s[0]*s[0], s[0]*s[1], s[0]*s[2],
+    s[1]*s[0], s[1]*s[1], s[1]*s[2],
+    s[2]*s[0], s[2]*s[1], s[2]*s[2]
+  };
+
+  // squared norm, norm of s
+  mjtNum xx = mju_dot3(s, s);
+  mjtNum x = mju_sqrt(xx);
+
+  // 4 coefficients: a=cos(x), b=sin(x)/x, c=(1-cos(x))/x^2, d=(x-sin(x))/x^3}
+  mjtNum a = mju_cos(x);
+  mjtNum b, c, d;
+
+  // x is not small: use full expressions
+  if (mju_abs(x) > 1.0/32) {
+    b = mju_sin(x) / x;
+    c = (1.0 - a) / xx;
+    d = (1.0 - b) / xx;
+  }
+
+  // |x| <= 1/32: use 6th order Taylor expansion (Horner form)
+  else {
+    b =  1 + xx/6  * (xx/20 * (1 - xx/42) - 1);
+    c = (1 + xx/12 * (xx/30 * (1 - xx/56) - 1)) / 2;
+    d = (1 + xx/20 * (xx/42 * (1 - xx/72) - 1)) / 6;
+  }
+
+  // derivatives
+  mjtNum Dvel_[9];
+  for (int i=0; i < 9; i++) {
+    if (Dquat)          Dquat[i] = a*eye[i] + b*cross[i] + c*outer[i];
+    if (Dvel || Dscale) Dvel_[i] = b*eye[i] + c*cross[i] + d*outer[i];
+  }
+  if (Dvel) mju_copy(Dvel, Dvel_, 9);
+  if (Dscale) mju_rotVecMat(Dscale, vel, Dvel_);
+}
+
+
+
 //------------------------- dense derivatives of component functions -------------------------------
 // no longer used, except in tests
 

src/engine/engine_derivative.h

@@ -23,9 +23,13 @@
 extern "C" {
 #endif
 
-// derivative of mju_subQuat w.r.t inputs
+// derivatives of mju_subQuat w.r.t inputs
 MJAPI void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db[9]);
 
+// derivatives of mju_quatIntegrate w.r.t inputs
+MJAPI void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
+                             mjtNum Dquat[9], mjtNum Dvel[9], mjtNum Dscale[3]);
+
 // analytical derivative of smooth forces w.r.t velocities:
 //   d->qDeriv = d (qfrc_actuator + qfrc_passive - [qfrc_bias]) / d qvel
 MJAPI void mjd_smooth_vel(const mjModel* m, mjData* d, int flg_bias);

test/engine/engine_derivative_test.cc

@@ -70,9 +70,9 @@ static mjtNum CompareMatrices(mjtNum* Actual, mjtNum* Expected,
   return max_error;
 }
 
-// utility function for matrix printing
+// utility function for matrix printing (debug)
+// NOLINTNEXTLINE(clang-diagnostic-unused-function)
 static void PrintMatrix(mjtNum* mat, int nrow, int ncol) {
-  // NOLINT(clang-diagnostic-unused-function)
   std::cerr.precision(5);
   std::cerr << "\n";
   for (int r=0; r < nrow; r++) {
@@ -836,5 +836,117 @@ TEST_F(DerivativeTest, SubQuat) {
   }
 }
 
+// utility: random quaternion, 3D velocity
+void randomQuatVel(mjtNum quat[4], mjtNum vel[3], int seed) {
+  // make distribution using seed
+  std::mt19937_64 rng;
+  rng.seed(seed);
+  std::normal_distribution<double> dist(0, 1);
+
+  // sample quat
+  for (int i=0; i < 4; i++) {
+    quat[i] = dist(rng);
+  }
+  mju_normalize4(quat);
+
+  // sample vel
+  for (int i=0; i < 3; i++) {
+    vel[i] = dist(rng);
+  }
+}
+
+// utility: finite-difference Jacobians of mju_quatIntegrate
+void mjd_quatIntegrateFD(mjtNum Dquat[9], mjtNum Ds[9],
+                         mjtNum Dvel[9], mjtNum Dh[3],
+                         const mjtNum quat[4], const mjtNum vel[3],
+                         mjtNum h, mjtNum eps) {
+  // compute y, output of mju_quatIntegrate(quat, vel, h)
+  mjtNum y[4] = {quat[0], quat[1], quat[2], quat[3]};
+  mju_quatIntegrate(y, vel, h);
+
+  mjtNum dx[3];      // nudged tangent-space input
+  mjtNum dq[4];      // quat output
+  mjtNum dy[3];      // nudged tangent-space output
+  mjtNum DquatT[9];  // Dquat transposed
+  mjtNum DsT[9];     // Ds transposed
+  mjtNum DvelT[9];   // Dvel transposed
+
+  for (int i = 0; i < 3; i++) {
+    // perturbation
+    mju_zero3(dx);
+    dx[i] = 1.0;
+
+    // d_y / d_quat
+    mju_copy4(dq, quat);
+    mju_quatIntegrate(dq, dx, eps);  // nudge dq
+    mju_quatIntegrate(dq, vel, h);  // compute nudged
+    mju_subQuat(dy, dq, y);          // subtract
+    mju_scl3(DquatT + i * 3, dy, 1.0 / eps);
+
+    // d_y / d_sv (scaled velocity)
+    mju_copy4(dq, quat);
+    mjtNum dsv[3] = {vel[0]*h, vel[1]*h, vel[2]*h};
+    mju_addToScl3(dsv, dx, eps);      // nudge dsv
+    mju_quatIntegrate(dq, dsv, 1.0);  // compute nudged
+    mju_subQuat(dy, dq, y);           // subtract
+    mju_scl3(DsT + i * 3, dy, 1.0 / eps);
+
+    // d_y / d_v (unscaled velocity)
+    mju_copy4(dq, quat);
+    mjtNum dv[3] = {vel[0], vel[1], vel[2]};
+    mju_addToScl3(dv, dx, eps);    // nudge dv
+    mju_quatIntegrate(dq, dv, h);  // compute nudged
+    mju_subQuat(dy, dq, y);        // subtract
+    mju_scl3(DvelT + i * 3, dy, 1.0 / eps);
+  }
+
+  // d_y / d_h (unscaled velocity)
+  mju_copy4(dq, quat);
+  mju_quatIntegrate(dq, vel, h + eps);  // compute nudged
+  mju_subQuat(dy, dq, y);               // subtract
+  mju_scl3(Dh, dy, 1.0 / eps);
+
+  // transpose
+  mju_transpose(Dquat, DquatT, 3, 3);
+  mju_transpose(Ds, DsT, 3, 3);
+  mju_transpose(Dvel, DsT, 3, 3);
+}
+
+TEST_F(DerivativeTest, quatIntegrate) {
+  const int nrepeats = 10;  // number of repeats
+  const mjtNum eps = 1e-7;  // epsilon for finite-differencing and comparison
+
+  int seed = 1;
+  for (int i = 0; i < nrepeats; i++) {
+    for (mjtNum h : {0.0, 1e-9, 1e-5, 1e-2, 1.0, 4.0}) {
+      // make random quaternion and velocity
+      mjtNum quat[4];
+      mjtNum vel[3];
+      randomQuatVel(quat, vel, seed++);
+
+      // analytic Jacobians
+      mjtNum Dquat[9];  // d_quatIntegrate(quat, vel, h) / d_quat
+      mjtNum Dvel[9];   // d_quatIntegrate(quat, vel, h) / d_vel
+      mjtNum Dh[3];     // d_quatIntegrate(quat, vel, h) / d_h
+      mjd_quatIntegrate(vel, h, Dquat, Dvel, Dh);
+
+      // finite-differenced Jacobians
+      mjtNum DquatFD[9];
+      mjtNum DsFD[9];
+      mjtNum DvelFD[9];
+      mjtNum DhFD[3];
+      mjd_quatIntegrateFD(DquatFD, DsFD, DvelFD, DhFD, quat, vel, h, eps);
+
+      // expect numerical equality of un/scaled velocity derivatives
+      EXPECT_THAT(AsVector(DvelFD, 9), Pointwise(DoubleNear(eps), DsFD));
+
+      // expect numerical equality of analytic and FD derivatives
+      EXPECT_THAT(AsVector(DquatFD, 9), Pointwise(DoubleNear(eps), Dquat));
+      EXPECT_THAT(AsVector(DvelFD, 9), Pointwise(DoubleNear(eps), Dvel));
+      EXPECT_THAT(AsVector(DhFD, 3), Pointwise(DoubleNear(eps), Dh));
+    }
+  }
+}
+
 }  // namespace
 }  // namespace mujoco

unity/Runtime/Bindings/MjBindings.cs

@@ -3865,5 +3865,11 @@ public static unsafe extern void mjd_transitionFD(mjModel_* m, mjData_* d, doubl
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mjd_inverseFD(mjModel_* m, mjData_* d, double eps, byte flg_actuation, double* DfDq, double* DfDv, double* DfDa, double* DsDq, double* DsDv, double* DsDa, double* DmDq);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mjd_subQuat(double* qa, double* qb, double* Da, double* Db);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mjd_quatIntegrate(double* vel, double scale, double* Dquat, double* Dvel, double* Dscale);
 }
 }