Merge pull request #1450 from v-r-a:angmomMat

PiperOrigin-RevId: 612835207
Change-Id: I8e7b01cc433563556a3ae4bbdd0c53612cc15e97

doc/APIreference/functions.rst

@@ -307,6 +307,18 @@ mj_jacPointAxis
 
 Compute translation end-effector Jacobian of point, and rotation Jacobian of axis.
 
+.. _mj_angmomMat:
+
+mj_angmomMat
+~~~~~~~~~~~~
+
+.. mujoco-include:: mj_angmomMat
+
+This function computes the ``3 x nv`` angular momentum matrix :math:`H(q)`, providing the linear mapping from
+generalized velocities to subtree angular momentum. More precisely if :math:`h` is the subtree angular momentum of
+body index ``body`` in ``mjData.subtree_angmom`` (reported by the :ref:`subtreeangmom<sensor-subtreeangmom>` sensor)
+and :math:`\dot q` is the generalized velocity ``mjData.qvel``, then :math:`h = H \dot q`.
+
 .. _mj_name2id:
 
 mj_name2id

doc/APIreference/functions_override.rst

@@ -173,6 +173,13 @@ the advantages of working in minimal coordinates.
 This and the remaining variants of the Jacobian function call mj_jac internally, with the center of the body, geom or
 site. They are just shortcuts; the same can be achieved by calling mj_jac directly.
 
+.. _mj_angmomMat:
+
+This function computes the ``3 x nv`` angular momentum matrix :math:`H(q)`, providing the linear mapping from
+generalized velocities to subtree angular momentum. More precisely if :math:`h` is the subtree angular momentum of
+body index ``body`` in ``mjData.subtree_angmom`` (reported by the :ref:`subtreeangmom<sensor-subtreeangmom>` sensor)
+and :math:`\dot q` is the generalized velocity ``mjData.qvel``, then :math:`h = H \dot q`.
+
 .. _mj_mulM:
 
 This function multiplies the joint-space inertia matrix stored in mjData.qM by a vector. qM has a custom sparse format

doc/changelog.rst

@@ -16,27 +16,30 @@ General
    its internal file buffers. :ref:`mj_addBufferVFS` allocates an empty buffer with a given name in an mjVFS and
    copies the data buffer into it, combining and replacing the deprecated two-step process of calling
    :ref:`mj_makeEmptyFileVFS` followed by a direct copy into the given mjVFS internal file buffer.
+3. Added :ref:`mj_angmomMat` which computes the ``3 x nv`` angular momentum matrix :math:`H(q)`, providing the linear
+   mapping from generalized velocities to subtree angular momentum :math:`h = H \dot q`. Contribution by
+   :github:user:`v-r-a`.
 
 MJX
 ^^^
 
-3. Improved performance of getting and putting device data.
+4. Improved performance of getting and putting device data.
 
    - Use ``tobytes()`` for numpy array serialization, which is orders of magnitude faster than converting to tuples.
    - Avoid reallocating host ``mjData`` arrays when array shapes are unchanged.
    - Speed up calculation of ``mjx.ncon`` for models with many geoms.
    - Avoid calling ``mjx.ncon`` in ``mjx.get_data_into`` when ``nc`` can be derived from ``mjx.Data``.
-4. Fixed a bug in ``mjx-viewer`` that prevented it from running.  Updated ``mjx-viewer`` to use newer
+5. Fixed a bug in ``mjx-viewer`` that prevented it from running.  Updated ``mjx-viewer`` to use newer
    ``mjx.get_data_into`` function call.
-5. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
+6. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
-6. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
+7. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
-7. Added support for :ref:`mjOption.impratio<mjOption>` to ``mjx.Model``.
+8. Added support for :ref:`mjOption.impratio<mjOption>` to ``mjx.Model``.
-8. Added support for cameras in ``mjx.Model`` and ``mjx.Data``. Fixes :github:issue:`1422`.
+9. Added support for cameras in ``mjx.Model`` and ``mjx.Data``. Fixes :github:issue:`1422`.
-9. Added an implementation of broadphase using `top_k` and bounding spheres.
+10. Added an implementation of broadphase using `top_k` and bounding spheres.
 
 Python bindings
 ^^^^^^^^^^^^^^^
-10. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members
+11. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members
     of the ``mjContact`` struct, and all array members of the ``mjrContext`` struct.
 
 

doc/includes/references.h

@@ -2517,6 +2517,7 @@ void mj_jacGeom(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, i
 void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* jacr, int site);
 void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
                      const mjtNum point[3], const mjtNum axis[3], int body);
+void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
 int mj_name2id(const mjModel* m, int type, const char* name);
 const char* mj_id2name(const mjModel* m, int type, int id);
 void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);

include/mujoco/mujoco.h

@@ -430,6 +430,9 @@ MJAPI void mj_jacSite(const mjModel* m, const mjData* d, mjtNum* jacp, mjtNum* j
 MJAPI void mj_jacPointAxis(const mjModel* m, mjData* d, mjtNum* jacPoint, mjtNum* jacAxis,
                            const mjtNum point[3], const mjtNum axis[3], int body);
 
+// Compute subtree angular momentum matrix.
+MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
+
 // Get id of object with the specified mjtObj type and name, returns -1 if id not found.
 MJAPI int mj_name2id(const mjModel* m, int type, const char* name);
 

introspect/functions.py

@@ -2402,6 +2402,36 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Compute translation end-effector Jacobian of point, and rotation Jacobian of axis.',  # pylint: disable=line-too-long
      )),
+    ('mj_angmomMat',
+     FunctionDecl(
+         name='mj_angmomMat',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='d',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='mat',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='body',
+                 type=ValueType(name='int'),
+             ),
+         ),
+         doc='Compute subtree angular momentum matrix.',
+     )),
     ('mj_name2id',
      FunctionDecl(
          name='mj_name2id',

python/mujoco/bindings_test.py

@@ -744,6 +744,16 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     # Expect next states to be equal.
     np.testing.assert_array_equal(state1a, state1b)
 
+  def test_mj_angmomMat(self):  # pylint: disable=invalid-name
+    self.data.qvel = np.ones(self.model.nv, np.float64)
+    mujoco.mj_forward(self.model, self.data)
+    mujoco.mj_subtreeVel(self.model, self.data)
+
+    mat = np.empty((3, 10), np.float64)
+    mujoco.mj_angmomMat(self.model, self.data, mat, 0)
+    np.testing.assert_almost_equal(mat @ self.data.qvel,
+                                   self.data.subtree_angmom[0, :])
+
   def test_mj_jacSite(self):  # pylint: disable=invalid-name
     mujoco.mj_forward(self.model, self.data)
     site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite')

python/mujoco/functions.cc

@@ -444,6 +444,14 @@ PYBIND11_MODULE(_functions, pymodule) {
             jacr.has_value() ? jacr->data() : nullptr,
             &(*point)[0], &(*axis)[0], body);
       });
+  Def<traits::mj_angmomMat>(
+      pymodule, [](const raw::MjModel* m, raw::MjData* d,
+                   Eigen::Ref<EigenArrayXX> mat, int body) {
+        if (mat.rows() != 3 || mat.cols() != m->nv) {
+          throw py::type_error("mat should be of shape (3, nv)");
+        }
+        return InterceptMjErrors(::mj_angmomMat)(m, d, mat.data(), body);
+      });
   Def<traits::mj_name2id>(pymodule);
   Def<traits::mj_id2name>(pymodule);
   Def<traits::mj_fullM>(

src/engine/engine_support.c

@@ -776,6 +776,76 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
 
 
 
+// compute subtree angular momentum matrix
+void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body) {
+  int nv = m->nv;
+  mj_markStack(d);
+
+  // stack allocations
+  mjtNum* jacp = mj_stackAllocNum(d, 3*nv);
+  mjtNum* jacr = mj_stackAllocNum(d, 3*nv);
+  mjtNum* term1 = mj_stackAllocNum(d, 3*nv);
+  mjtNum* term2 = mj_stackAllocNum(d, 3*nv);
+
+  // clear output
+  mju_zero(mat, 3*nv);
+
+  // save the location of the subtree COM
+  mjtNum subtree_com[3];
+  mju_copy3(subtree_com, d->subtree_com+3*body);
+
+  for (int b=body; b < m->nbody; b++) {
+    // end of body subtree, break from the loop
+    if (b > body && m->body_parentid[b] < body) {
+      break;
+    }
+
+    // linear and angular velocity Jacobian of the body COM (inertial frame)
+    mj_jacBodyCom(m, d, jacp, jacr, b);
+
+    // orientation of the COM (inertial) frame of b-th body
+    mjtNum ximat[9];
+    mju_copy(ximat, d->ximat+9*b, 9);
+
+    // save the inertia matrix of b-th body
+    mjtNum inertia[9] = {0};
+    inertia[0] = m->body_inertia[3*b];   // inertia(1,1)
+    inertia[4] = m->body_inertia[3*b+1]; // inertia(2,2)
+    inertia[8] = m->body_inertia[3*b+2]; // inertia(3,3)
+
+    // term1 = body angular momentum about self COM in world frame
+    mjtNum tmp1[9], tmp2[9];
+    mju_mulMatMat(tmp1, ximat, inertia, 3, 3, 3);  // tmp1  = ximat * inertia
+    mju_mulMatMatT(tmp2, tmp1, ximat, 3, 3, 3);    // tmp2  = ximat * inertia * ximat^T
+    mju_mulMatMat(term1, tmp2, jacr, 3, 3, nv);    // term1 = ximat * inertia * ximat^T * jacr
+
+    // location of body COM w.r.t subtree COM
+    mjtNum com[3];
+    mju_sub3(com, d->xipos+3*b, subtree_com);
+
+    // skew symmetric matrix representing body_com vector
+    mjtNum com_mat[9] = {0};
+    com_mat[1] = -com[2];
+    com_mat[2] = com[1];
+    com_mat[3] = com[2];
+    com_mat[5] = -com[0];
+    com_mat[6] = -com[1];
+    com_mat[7] = com[0];
+
+    // term2 = moment of linear momentum
+    mju_mulMatMat(term2, com_mat, jacp, 3, 3, nv);   // term2 = com_mat * jacp
+    mju_scl(term2, term2, m->body_mass[b], 3 * nv);  // term2 = com_mat * jacp * mass
+
+    // mat += term1 + term2
+    mju_addTo(mat, term1, 3*nv);
+    mju_addTo(mat, term2, 3*nv);
+  }
+
+  mj_freeStack(d);
+}
+
+
+
 //-------------------------- name functions --------------------------------------------------------
 
 // get number of objects and name addresses for given object type

src/engine/engine_support.h

@@ -107,6 +107,9 @@ int mj_jacSum(const mjModel* m, mjData* d, int* chain,
               int n, const int* body, const mjtNum* weight,
               const mjtNum point[3], mjtNum* jac, int flg_rot);
 
+// compute subtree angular momentum matrix
+MJAPI void mj_angmomMat(const mjModel* m, mjData* d, mjtNum* mat, int body);
+
 
 //-------------------------- name functions --------------------------------------------------------
 

test/engine/engine_support_test.cc

@@ -16,6 +16,7 @@
 
 #include "src/engine/engine_support.h"
 
+#include <limits>
 #include <random>
 #include <string>
 #include <vector>
@@ -37,6 +38,125 @@ using ::testing::ContainsRegex;
 using ::testing::MatchesRegex;
 using ::testing::Pointwise;
 using ::testing::ElementsAreArray;
+
+using AngMomMatTest = MujocoTest;
+
+static constexpr char AngMomTestingModel[] = R"(
+<mujoco>
+  <option>
+    <flag gravity="disable"/>
+  </option>
+  <worldbody>
+    <body name="link1" pos="0 0 0.5">
+      <freejoint/>
+      <geom type="ellipsoid" size="0.15 0.17 0.19" quat="1 .2 .3 .4"/>
+      <body name="link2" >
+        <joint type="hinge" axis="1 0 0" />
+        <geom type="capsule" size="0.05" fromto="0 0 0  0 0.5 0"/>
+        <body pos="0 0.6 0">
+          <joint type="slide" axis="1 0 0"/>
+          <geom type="capsule" size="0.05 0.2" quat="0.707 0 0.707 0"/>
+          <body name="link3">
+            <joint type="ball" pos="0.2 0 0"/>
+            <geom type="capsule" pos="0.2 0 0" size="0.03 0.4"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <keyframe>
+    <key qvel="0 0 0 .1 .2 .3 .4 .5 .4 .3 .2"/>
+  </keyframe>
+</mujoco>
+)";
+
+// compare subtree angular momentum computed in two ways
+TEST_F(AngMomMatTest, CompareAngMom) {
+  mjModel* model = LoadModelFromString(AngMomTestingModel);
+  int nv = model->nv;
+  int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
+
+  mjData* data = mj_makeData(model);
+
+  // reset to the keyframe with some angular velocities
+  mj_resetDataKeyframe(model, data, 0);
+  mj_forward(model, data);
+
+  // get the reference value of angular momentum
+  mj_subtreeVel(model, data);
+  mjtNum angmom_ref[3];
+  mju_copy3(angmom_ref, data->subtree_angmom+3*bodyid);
+
+  // compute angular momentum using the angular momentum matrix
+  mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+  mj_angmomMat(model, data, angmom_mat, bodyid);
+  mjtNum angmom_test[3];
+  mju_mulMatVec(angmom_test, angmom_mat, data->qvel, 3, nv);
+
+  // compare the two angular momentum values
+  static const mjtNum tol = 1e-8;
+  for (int i = 0; i < 3; i++) {
+    EXPECT_THAT(angmom_ref[i], DoubleNear(angmom_test[i], tol));
+  }
+
+  mju_free(angmom_mat);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+// compare subtree angular momentum matrix: analytical and findiff
+TEST_F(AngMomMatTest, CompareAngMomMats) {
+  mjModel* model = LoadModelFromString(AngMomTestingModel);
+  int nv = model->nv;
+  int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
+  mjData* data = mj_makeData(model);
+  mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+  mjtNum* angmom_mat_fd = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+
+  // reset to the keyframe with some angular velocities
+  mj_resetDataKeyframe(model, data, 0);
+  mj_forward(model, data);
+
+  // compute the angular momentum matrix using the analytical method
+  mj_angmomMat(model, data, angmom_mat, bodyid);
+
+  // compute the angular momentum matrix using finite differences
+  static const mjtNum eps = 1e-6;
+  for (int i = 0; i < nv; i++) {
+    // reset vel, forward nudge i-th dof, get angmom
+    mju_copy(data->qvel, model->key_qvel, model->nv);
+    data->qvel[i] += eps;
+    mj_forward(model, data);
+    mj_subtreeVel(model, data);
+    mjtNum agmf[3];
+    mju_copy3(agmf, data->subtree_angmom+3*bodyid);
+
+    // reset vel, backward nudge i-th dof, get angmom
+    mju_copy(data->qvel, model->key_qvel, model->nv);
+    data->qvel[i] -= eps;
+    mj_forward(model, data);
+    mj_subtreeVel(model, data);
+    mjtNum agmb[3];
+    mju_copy3(agmb, data->subtree_angmom+3*bodyid);
+
+    // finite-difference the angmom matrix
+    for (int j = 0; j < 3; j++) {
+      angmom_mat_fd[nv*j+i] = (agmf[j] - agmb[j]) / (2 * eps);
+    }
+  }
+
+  // compare the two matrices
+  static const mjtNum tol = 1e-8;
+  for (int i = 0; i < 3*nv; i++) {
+    EXPECT_THAT(angmom_mat_fd[i], DoubleNear(angmom_mat[i], tol));
+  }
+
+  mju_free(angmom_mat_fd);
+  mju_free(angmom_mat);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 using JacobianTest = MujocoTest;
 static const mjtNum max_abs_err = std::numeric_limits<float>::epsilon();
 
@@ -446,7 +566,7 @@ TEST_F(AddMTest, DenseSameAsSparse) {
           rowadr.data(), colind.data());
 
   // dense addM
-  mj_addM(m, d, dst_dense.data(), NULL, NULL, NULL);
+  mj_addM(m, d, dst_dense.data(), nullptr, nullptr, nullptr);
 
   // dense comparison, should be same matrix
   EXPECT_THAT(dst_dense, ElementsAreArray(dst_sparse));

unity/Runtime/Bindings/MjBindings.cs

@@ -6627,6 +6627,9 @@ public static unsafe extern void mj_jacSite(mjModel_* m, mjData_* d, double* jac
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_jacPointAxis(mjModel_* m, mjData_* d, double* jacPoint, double* jacAxis, double* point, double* axis, int body);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mj_angmomMat(mjModel_* m, mjData_* d, double* mat, int body);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern int mj_name2id(mjModel_* m, int type, [MarshalAs(UnmanagedType.LPStr)]string name);