Add mju_mat2Rot.

This function extracts the 3D rotation from an arbitrary 3x3 matrix by refining the input quaternion. It is based on the paper "A robust method to extract the rotational part of deformations" by Müller, Matthias, Jan Bender, Nuttapong Chentanez, and Miles Macklin.

PiperOrigin-RevId: 700006006
Change-Id: I77550993233dea9cdf68601762a3ae7ded749bdf

doc/APIreference/functions.rst

@@ -3555,6 +3555,16 @@ Integrate quaternion given 3D angular velocity.
 
 Construct quaternion performing rotation from z-axis to given vector.
 
+.. _mju_mat2Rot:
+
+`mju_mat2Rot <#mju_mat2Rot>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mju_mat2Rot
+
+extract 3D rotation from an arbitrary 3x3 matrix by refining the input quaternion
+returns the number of iterations required to converge
+
 .. _mju_euler2Quat:
 
 `mju_euler2Quat <#mju_euler2Quat>`__

doc/includes/references.h

@@ -3481,6 +3481,7 @@ void mju_mat2Quat(mjtNum quat[4], const mjtNum mat[9]);
 void mju_derivQuat(mjtNum res[4], const mjtNum quat[4], const mjtNum vel[3]);
 void mju_quatIntegrate(mjtNum quat[4], const mjtNum vel[3], mjtNum scale);
 void mju_quatZ2Vec(mjtNum quat[4], const mjtNum vec[3]);
+int mju_mat2Rot(mjtNum quat[4], const mjtNum mat[9]);
 void mju_euler2Quat(mjtNum quat[4], const mjtNum euler[3], const char* seq);
 void mju_mulPose(mjtNum posres[3], mjtNum quatres[4],
                  const mjtNum pos1[3], const mjtNum quat1[4],

include/mujoco/mujoco.h

@@ -1120,6 +1120,10 @@ MJAPI void mju_quatIntegrate(mjtNum quat[4], const mjtNum vel[3], mjtNum scale);
 // Construct quaternion performing rotation from z-axis to given vector.
 MJAPI void mju_quatZ2Vec(mjtNum quat[4], const mjtNum vec[3]);
 
+// extract 3D rotation from an arbitrary 3x3 matrix by refining the input quaternion
+// returns the number of iterations required to converge
+MJAPI int mju_mat2Rot(mjtNum quat[4], const mjtNum mat[9]);
+
 // Convert sequence of Euler angles (radians) to quaternion.
 // seq[0,1,2] must be in 'xyzXYZ', lower/upper-case mean intrinsic/extrinsic rotations.
 MJAPI void mju_euler2Quat(mjtNum quat[4], const mjtNum euler[3], const char* seq);

introspect/functions.py

@@ -7385,6 +7385,28 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Construct quaternion performing rotation from z-axis to given vector.',  # pylint: disable=line-too-long
      )),
+    ('mju_mat2Rot',
+     FunctionDecl(
+         name='mju_mat2Rot',
+         return_type=ValueType(name='int'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='quat',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(4,),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='mat',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                     extents=(9,),
+                 ),
+             ),
+         ),
+         doc='extract 3D rotation from an arbitrary 3x3 matrix by refining the input quaternion returns the number of iterations required to converge',  # pylint: disable=line-too-long
+     )),
     ('mju_euler2Quat',
      FunctionDecl(
          name='mju_euler2Quat',

python/mujoco/functions.cc

@@ -1069,6 +1069,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mju_derivQuat>(pymodule);
   Def<traits::mju_quatIntegrate>(pymodule);
   Def<traits::mju_quatZ2Vec>(pymodule);
+  Def<traits::mju_mat2Rot>(pymodule);
   Def<traits::mju_euler2Quat>(pymodule);
 
   // Poses

src/engine/engine_util_spatial.c

@@ -286,6 +286,47 @@ void mju_quatZ2Vec(mjtNum quat[4], const mjtNum vec[3]) {
 
 
 
+// extract 3D rotation from an arbitrary 3x3 matrix
+static const mjtNum rotEPS = 1e-9;
+int mju_mat2Rot(mjtNum quat[4], const mjtNum mat[9]) {
+  // Müller, Matthias, Jan Bender, Nuttapong Chentanez, and Miles Macklin. "A
+  // robust method to extract the rotational part of deformations." In
+  // Proceedings of the 9th International Conference on Motion in Games, pp.
+  // 55-60. 2016.
+
+  int iter;
+  mjtNum col1_mat[3] = {mat[0], mat[3], mat[6]};
+  mjtNum col2_mat[3] = {mat[1], mat[4], mat[7]};
+  mjtNum col3_mat[3] = {mat[2], mat[5], mat[8]};
+  for (iter = 0; iter < 500; iter++) {
+    mjtNum rot[9];
+    mju_quat2Mat(rot, quat);
+    mjtNum col1_rot[3] = {rot[0], rot[3], rot[6]};
+    mjtNum col2_rot[3] = {rot[1], rot[4], rot[7]};
+    mjtNum col3_rot[3] = {rot[2], rot[5], rot[8]};
+    mjtNum omega[3], vec1[3], vec2[3], vec3[3];
+    mju_cross(vec1, col1_rot, col1_mat);
+    mju_cross(vec2, col2_rot, col2_mat);
+    mju_cross(vec3, col3_rot, col3_mat);
+    mju_add3(omega, vec1, vec2);
+    mju_addTo3(omega, vec3);
+    mju_scl3(omega, omega, 1.0 / (mju_abs(mju_dot3(col1_rot, col1_mat) +
+                                          mju_dot3(col2_rot, col2_mat) +
+                                          mju_dot3(col3_rot, col3_mat)) + mjMINVAL));
+    mjtNum w = mju_normalize3(omega);
+    if (w < rotEPS) {
+      break;
+    }
+    mjtNum qrot[4];
+    mju_axisAngle2Quat(qrot, omega, w);
+    mju_mulQuat(quat, qrot, quat);
+    mju_normalize4(quat);
+  }
+  return iter;
+}
+
+
+
 //------------------------------ pose operations (quat, pos) ---------------------------------------
 
 // multiply two poses

src/engine/engine_util_spatial.h

@@ -59,6 +59,10 @@ MJAPI void mju_quatIntegrate(mjtNum quat[4], const mjtNum vel[3], mjtNum scale);
 // compute quaternion performing rotation from z-axis to given vector
 MJAPI void mju_quatZ2Vec(mjtNum quat[4], const mjtNum vec[3]);
 
+// extract 3D rotation from an arbitrary 3x3 matrix by refining the input quaternion
+// returns the number of iterations required to converge
+MJAPI int mju_mat2Rot(mjtNum quat[4], const mjtNum mat[9]);
+
 
 //------------------------------ pose operations (pos, quat) ---------------------------------------
 

test/engine/engine_util_spatial_test.cc

@@ -15,6 +15,7 @@
 // Tests for engine/engine_util_spatial.c
 
 #include <cmath>
+#include <cstdlib>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -196,5 +197,69 @@ TEST_F(Euler2QuatTest, Euler2Quat) {
   EXPECT_THAT(quat, Pointwise(DoubleNear(tol), expected6));
 }
 
+using Mat2RotTest = MujocoTest;
+
+TEST_F(Mat2RotTest, RotationFromArbitraryMatrix) {
+  // create arbitrary target rotation matrix
+  mjtNum target[4], rot[9];
+  mjtNum axis[3] = {1, 1, 1};
+  mju_axisAngle2Quat(target, axis, mjPI/6);
+  mju_normalize4(target);
+  mju_quat2Mat(rot, target);
+
+  // combine rotation with arbitrary stretch
+  mjtNum mat[9];
+  mjtNum deformation_gradient[9] = {0.5, 0.25, 0.125,
+                                    0.3, 0.66, 0.999,
+                                    0.4, 0.22, 0.111};
+  mjtNum stretch[9];
+  mju_mulMatTMat3(stretch, deformation_gradient, deformation_gradient);
+  mju_mulMatMat3(mat, rot, stretch);
+
+  // calculate rotational part of the matrix
+  mjtNum quat[4] = {1, 0, 0, 0};
+  int niter = mju_mat2Rot(quat, mat);
+  EXPECT_THAT(quat, Pointwise(DoubleNear(1e-8), target));
+  EXPECT_LE(niter, 150);
+}
+
+TEST_F(Mat2RotTest, IdentityFromRandomRotation) {
+  // This test is based on the following paper:
+  // Müller, Matthias, Jan Bender, Nuttapong Chentanez, and Miles Macklin. "A
+  // robust method to extract the rotational part of deformations." In
+  // Proceedings of the 9th International Conference on Motion in Games, pp.
+  // 55-60. 2016.
+  mjtNum mat[9] = {1, 0, 0, 0, 1, 0, 0, 0, 1};
+  srand(123);
+
+  for (int i = 0; i < 100; ++i) {
+    // random quaternion
+    mjtNum quat[4];
+    for (int j = 0; j < 4; ++j) {
+      quat[j] = rand() / (float)RAND_MAX;  // NOLINT
+    }
+
+    // calculate rotational part of the matrix
+    mjtNum res[9];
+    mju_normalize4(quat);
+    EXPECT_LE(mju_mat2Rot(quat, mat), 40);
+    mju_quat2Mat(res, quat);
+    EXPECT_THAT(res, Pointwise(DoubleNear(1e-6), mat));
+  }
+}
+
+TEST_F(Mat2RotTest, SpecialCases) {
+  mjtNum eye[9] = {1, 0, 0, 0, 1, 0, 0, 0, 1};
+  mjtNum quat[4] = {1, 0, 0, 0};
+  EXPECT_EQ(mju_mat2Rot(quat, eye), 0);
+  EXPECT_THAT(quat, Pointwise(DoubleNear(1e-8), {1, 0, 0, 0}));
+  mjtNum zero[9] = {0, 0, 0, 0, 0, 0, 0, 0, 0};
+  EXPECT_EQ(mju_mat2Rot(quat, zero), 0);
+  EXPECT_THAT(quat, Pointwise(DoubleNear(1e-8), {1, 0, 0, 0}));
+  mjtNum ones[9] = {1, 1, 1, 1, 1, 1, 1, 1, 1};
+  EXPECT_EQ(mju_mat2Rot(quat, ones), 0);
+  EXPECT_THAT(quat, Pointwise(DoubleNear(1e-8), {1, 0, 0, 0}));
+}
+
 }  // namespace
 }  // namespace mujoco

unity/Runtime/Bindings/MjBindings.cs

@@ -7276,6 +7276,9 @@ public static unsafe extern void mju_quatIntegrate(double* quat, double* vel, do
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mju_quatZ2Vec(double* quat, double* vec);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern int mju_mat2Rot(double* quat, double* mat);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mju_euler2Quat(double* quat, double* euler, [MarshalAs(UnmanagedType.LPStr)]string seq);