diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index ce295ecc..c21e34f7 100644 --- a/doc/APIreference/functions.rst +++ b/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) +and :math:`\dot q` is the generalized velocity ``mjData.qvel``, then :math:`h = H \dot q`. + .. _mj_name2id: mj_name2id diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 812e6f43..ea7ec0cc 100644 --- a/doc/APIreference/functions_override.rst +++ b/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) +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 diff --git a/doc/changelog.rst b/doc/changelog.rst index bf1a4540..1f2e2b5b 100644 --- a/doc/changelog.rst +++ b/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.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`. -7. Added support for :ref:`mjOption.impratio` to ``mjx.Model``. -8. 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. +6. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices. +7. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`. +8. Added support for :ref:`mjOption.impratio` to ``mjx.Model``. +9. Added support for cameras in ``mjx.Model`` and ``mjx.Data``. Fixes :github:issue:`1422`. +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. diff --git a/doc/includes/references.h b/doc/includes/references.h index a61663d1..4c48acbd 100644 --- a/doc/includes/references.h +++ b/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); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 2e11f2ee..e08a911b 100644 --- a/include/mujoco/mujoco.h +++ b/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); diff --git a/introspect/functions.py b/introspect/functions.py index b2f7655f..5a211ffe 100644 --- a/introspect/functions.py +++ b/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', diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 9c377ecf..4a1da339 100644 --- a/python/mujoco/bindings_test.py +++ b/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') diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 4c7bed7c..db1a874f 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -444,6 +444,14 @@ PYBIND11_MODULE(_functions, pymodule) { jacr.has_value() ? jacr->data() : nullptr, &(*point)[0], &(*axis)[0], body); }); + Def( + pymodule, [](const raw::MjModel* m, raw::MjData* d, + Eigen::Ref 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(pymodule); Def(pymodule); Def( diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 38df9a14..c27c1374 100644 --- a/src/engine/engine_support.c +++ b/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 diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h index abb8a4ff..122c59ba 100644 --- a/src/engine/engine_support.h +++ b/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 -------------------------------------------------------- diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index a2b3acda..d6a892c1 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -16,6 +16,7 @@ #include "src/engine/engine_support.h" +#include #include #include #include @@ -37,6 +38,125 @@ using ::testing::ContainsRegex; using ::testing::MatchesRegex; using ::testing::Pointwise; using ::testing::ElementsAreArray; + +using AngMomMatTest = MujocoTest; + +static constexpr char AngMomTestingModel[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +// 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::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)); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index e8d6cabf..771c7659 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/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);