From 767c607f58b702e91f2050ba141da23dc067c70d Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 4 May 2026 14:25:23 -0700 Subject: [PATCH] Add `mju_sym2dense`, document future breakage of `mj_fullM` PiperOrigin-RevId: 910242375 Change-Id: Ibfbdef9cfb66088723499ea257da09aee0d80938 --- doc/APIreference/functions.rst | 9 ++++++ doc/changelog.rst | 14 +++++++++ doc/includes/references.h | 2 ++ include/mujoco/mujoco.h | 4 +++ mjx/mujoco/mjx/_src/io.py | 8 ++++- mjx/mujoco/mjx/_src/io_test.py | 6 ++-- mjx/mujoco/mjx/_src/smooth_test.py | 2 +- mjx/mujoco/mjx/_src/support_test.py | 2 +- mjx/mujoco/mjx/warp/forward_test.py | 2 +- python/mujoco/functions.cc | 31 +++++++++++++++++++ python/mujoco/introspect/functions.py | 42 ++++++++++++++++++++++++++ src/engine/engine_island.c | 2 +- src/engine/engine_setconst.c | 2 +- src/engine/engine_util_sparse.c | 17 +++++++++++ src/engine/engine_util_sparse.h | 4 +++ test/engine/engine_support_test.cc | 7 +++-- test/engine/engine_util_sparse_test.cc | 36 ++++++++++++++++++++++ unity/Runtime/Bindings/MjBindings.cs | 3 ++ wasm/codegen/generated/bindings.cc | 10 ++++++ 19 files changed, 191 insertions(+), 12 deletions(-) diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index fc513e20..a91344ac 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -3982,6 +3982,15 @@ Convert matrix from dense to sparse. Convert matrix from sparse to dense. +.. _mju_sym2dense: + +`mju_sym2dense <#mju_sym2dense>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mju_sym2dense + +Convert lower-triangular symmetric CSR matrix to full dense matrix. + .. _Quaternions: Quaternions diff --git a/doc/changelog.rst b/doc/changelog.rst index 26f1f591..92df67c5 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -5,6 +5,8 @@ Changelog Upcoming version (not yet released) ----------------------------------- +General +^^^^^^^ - Added island support for the :ref:`PGS solver`. - Added support for :ref:`elastic2d` for trilinear and quadratic flex :ref:`dofs`. @@ -13,6 +15,18 @@ Upcoming version (not yet released) (nonzero :ref:`density` or :ref:`viscosity`). Midpoint integration treats external forces as zero-order-hold constants, which causes energy gain in the presence of contacts and in fluid media. +- Added :ref:`mju_sym2dense`, converting a lower-triangular, implicitly symmetric CSR matrix to a dense + symmetric matrix. The inertia matrix ``mjData.M`` is an example of such a matrix. + +.. admonition:: Future breaking API changes + :class: warning + + - The introduction of :ref:`mju_sym2dense` is a step towards the removal of the legacy-format ``mjData.qM`` in favor + of the CSR-format ``mjData.M``. This removal will involve a future breaking change to :ref:`mj_fullM` (which + currently accepts a ``qM``-like matrix as an argument). To prevent a future breakage, replace + ``mj_fullM(m, dst, d->qM)`` with + |br| ``mju_sym2dense(dst, d->M, m->nv, m->M_rownnz, m->M_rowadr, m->M_colind)``. + Python ^^^^^^ diff --git a/doc/includes/references.h b/doc/includes/references.h index c594a32d..edc8b0c4 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3522,6 +3522,8 @@ int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind, int nnz); void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); +void mju_sym2dense(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind); void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]); void mju_negQuat(mjtNum res[4], const mjtNum quat[4]); void mju_mulQuat(mjtNum res[4], const mjtNum quat1[4], const mjtNum quat2[4]); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index c8ca6ef5..45b56bb1 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1206,6 +1206,10 @@ MJAPI int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc, MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); +// Convert lower-triangular symmetric CSR matrix to full dense matrix. +MJAPI void mju_sym2dense(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind); + //---------------------------------- Quaternions --------------------------------------------------- diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 918b0fa4..a4630703 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -1058,7 +1058,13 @@ def _put_data_jax( # convert qM and qLD if jacobian is dense if not support.is_sparse(m): impl_fields['qM'] = np.zeros((m.nv, m.nv)) - mujoco.mj_fullM(m, impl_fields['qM'], d.qM) + mujoco.mju_sym2dense( + impl_fields['qM'], + d.M, + m.M_rownnz, + m.M_rowadr, + m.M_colind, + ) # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing try: impl_fields['qLD'], _ = scipy.linalg.cho_factor(impl_fields['qM']) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 542f0718..6c61e8d6 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -520,7 +520,7 @@ class DataIOTest(parameterized.TestCase): if impl == 'jax': # check that qM is transformed properly qm = np.zeros((m.nv, m.nv), dtype=np.float64) - mujoco.mj_fullM(m, qm, d.qM) + mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind) np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx)) elif impl == 'cpp': @@ -529,7 +529,7 @@ class DataIOTest(parameterized.TestCase): return # cpp does not populate other fields in _impl elif impl == 'warp': qm = np.zeros((m.nv, m.nv), dtype=np.float64) - mujoco.mj_fullM(m, qm, d.qM) + mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind) np.testing.assert_allclose(dx._impl.qM, qm) # TODO(taylorhowell): test efc__J np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3]) @@ -596,7 +596,7 @@ class DataIOTest(parameterized.TestCase): dx_from_dense = mjx.put_data(m, d, impl=impl) if impl == 'jax': qm = np.zeros((m.nv, m.nv)) - mujoco.mj_fullM(m, qm, d.qM) + mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind) np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8) diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 7f0d8c09..b421b68f 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -438,7 +438,7 @@ class TendonTest(parameterized.TestCase): if jacobian == JacobianType.DENSE: qM = np.zeros((m.nv, m.nv)) # pylint: disable=invalid-name - mujoco.mj_fullM(m, qM, d.qM) + mujoco.mju_sym2dense(qM, d.M, m.M_rownnz, m.M_rowadr, m.M_colind) else: qM = d.qM # pylint: disable=invalid-name _assert_eq(dx._impl.qM, qM, 'qM') diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 6268f14b..a01fb004 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -64,7 +64,7 @@ class SupportTest(parameterized.TestCase): dx = mjx.put_data(m, d) mjx_full_m = jax.jit(support.full_m)(mx, dx) mj_full_m = np.zeros((m.nv, m.nv), dtype=np.float64) - mujoco.mj_fullM(m, mj_full_m, d.qM) + mujoco.mju_sym2dense(mj_full_m, d.M, m.M_rownnz, m.M_rowadr, m.M_colind) np.testing.assert_allclose(mjx_full_m, mj_full_m, atol=5e-5, rtol=5e-5) @parameterized.parameters('constraints.xml', 'pendula.xml') diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py index 056a0c51..2c38891e 100644 --- a/mjx/mujoco/mjx/warp/forward_test.py +++ b/mjx/mujoco/mjx/warp/forward_test.py @@ -175,7 +175,7 @@ class ForwardTest(parameterized.TestCase): tu.assert_attr_eq(dx._impl, d, 'crb') qm = np.zeros((m.nv, m.nv)) - mujoco.mj_fullM(m, qm, d.qM) + mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind) # mjwarp adds padding to qM tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM') # qLD is fused in a cholesky factorize and solve, and not written to. diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 03a1be1b..5787d756 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -1241,6 +1241,37 @@ PYBIND11_MODULE(_functions, pymodule) { colind.data()); }); + DEF_WITH_OMITTED_PY_ARGS(traits::mju_sym2dense, "n")( + pymodule, + [](Eigen::Ref res, + Eigen::Ref mat, + Eigen::Ref rownnz, + Eigen::Ref rowadr, + Eigen::Ref colind) { + if (res.rows() != res.cols()) { + throw py::type_error("res should be a square matrix"); + } + if (res.rows() != rownnz.size()) { + throw py::type_error("#rows in res should equal size of rownnz"); + } + if (res.rows() != rowadr.size()) { + throw py::type_error("#rows in res should equal size of rowadr"); + } + if (res.rows() > 0) { + int nnz = rowadr.array().tail(1)[0] + rownnz.array().tail(1)[0]; + if (mat.size() < nnz) { + throw py::type_error("mat size is too small for the given sparse " + "structure"); + } + if (colind.size() < nnz) { + throw py::type_error("colind size is too small for the given " + "sparse structure"); + } + } + return ::mju_sym2dense(res.data(), mat.data(), res.rows(), + rownnz.data(), rowadr.data(), colind.data()); + }); + // Quaternions Def(pymodule); Def(pymodule); diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py index e16f434f..ff12151e 100644 --- a/python/mujoco/introspect/functions.py +++ b/python/mujoco/introspect/functions.py @@ -7755,6 +7755,48 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Convert matrix from sparse to dense.', )), + ('mju_sym2dense', + FunctionDecl( + name='mju_sym2dense', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='res', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='mat', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), + FunctionParameterDecl( + name='n', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='rownnz', + type=PointerType( + inner_type=ValueType(name='int', is_const=True), + ), + ), + FunctionParameterDecl( + name='rowadr', + type=PointerType( + inner_type=ValueType(name='int', is_const=True), + ), + ), + FunctionParameterDecl( + name='colind', + type=PointerType( + inner_type=ValueType(name='int', is_const=True), + ), + ), + ), + doc='Convert lower-triangular symmetric CSR matrix to full dense matrix.', # pylint: disable=line-too-long + )), ('mju_rotVecQuat', FunctionDecl( name='mju_rotVecQuat', diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c index 399fcbc1..ebda019d 100644 --- a/src/engine/engine_island.c +++ b/src/engine/engine_island.c @@ -525,7 +525,7 @@ void mj_island(const mjModel* m, mjData* d) { d->island_dofadr[i] = d->map_idof2dof[d->island_idofadr[i]]; } - // inertia: block-diagonalize both iLD <- qLD and iM <- qM + // inertia: block-diagonalize both iLD <- qLD and iM <- M mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind, d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind, nidof, nisland, diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 754aa29b..e3e41069 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -1307,7 +1307,7 @@ static void setStat(mjModel* m, mjData* d) { if (m->nv) { m->stat.meaninertia = 0; for (int i=0; i < m->nv; i++) { - m->stat.meaninertia += d->qM[m->dof_Madr[i]]; + m->stat.meaninertia += d->M[m->M_rowadr[i] + m->M_rownnz[i] - 1]; } m->stat.meaninertia /= m->nv; } diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c index b4a74b1a..0c79a27f 100644 --- a/src/engine/engine_util_sparse.c +++ b/src/engine/engine_util_sparse.c @@ -140,6 +140,23 @@ void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, } +// convert lower-triangular symmetric CSR matrix to full dense matrix +void mju_sym2dense(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind) { + mju_zero(res, n*n); + for (int i = 0; i < n; i++) { + int adr = rowadr[i]; + for (int j = 0; j < rownnz[i]; j++) { + int col = colind[adr+j]; + if (col <= i) { + res[i*n+col] = mat[adr+j]; + res[col*n+i] = mat[adr+j]; + } + } + } +} + + // res[row, :] = mat[row, :] void mju_copySparse(mjtNum* res, const mjtNum* mat, const int* rownnz, const int* rowadr, const int* row, int nrow) { diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h index 3d5badaa..038ac5ce 100644 --- a/src/engine/engine_util_sparse.h +++ b/src/engine/engine_util_sparse.h @@ -46,6 +46,10 @@ MJAPI int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc, MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, const int* rownnz, const int* rowadr, const int* colind); +// convert lower-triangular symmetric CSR matrix to full dense matrix +MJAPI void mju_sym2dense(mjtNum* res, const mjtNum* mat, int n, + const int* rownnz, const int* rowadr, const int* colind); + // res[row, :] = mat[row, :] void mju_copySparse(mjtNum* res, const mjtNum* mat, const int* rownnz, const int* rowadr, const int* row, int nrow); diff --git a/test/engine/engine_support_test.cc b/test/engine/engine_support_test.cc index f52d374f..30bb5dfe 100644 --- a/test/engine/engine_support_test.cc +++ b/test/engine/engine_support_test.cc @@ -548,7 +548,8 @@ TEST_F(InertiaTest, mulM) { // dense M matrix vector Mdense(nv*nv); - mj_fullM(model, Mdense.data(), data->qM); + mju_sym2dense(Mdense.data(), data->M, nv, + model->M_rownnz, model->M_rowadr, model->M_colind); // arbitrary RHS vector vector vec(nv); @@ -611,9 +612,9 @@ TEST_F(InertiaTest, FullM) { mjData* d = mj_makeData(m); mj_forward(m, d); - // get dense mass matrix from qM using mj_fullM + // get dense mass matrix from M using mju_sym2dense vector M(nv * nv); - mj_fullM(m, M.data(), d->qM); + mju_sym2dense(M.data(), d->M, nv, m->M_rownnz, m->M_rowadr, m->M_colind); // get dense mass matrix from M using mju_sparse2dense vector M_CSR(nv * nv); diff --git a/test/engine/engine_util_sparse_test.cc b/test/engine/engine_util_sparse_test.cc index 67eadbbc..4e0dd3de 100644 --- a/test/engine/engine_util_sparse_test.cc +++ b/test/engine/engine_util_sparse_test.cc @@ -361,6 +361,42 @@ TEST_F(EngineUtilSparseTest, MjuCompressSparse) { EXPECT_EQ(AsVector(dense, 6), AsVector(dense_expected_minval1, 6)); } +TEST_F(EngineUtilSparseTest, MjuSym2Dense) { + // lower-triangular CSR for a 3x3 symmetric matrix: + // 1 2 0 + // 2 3 4 + // 0 4 5 + // stored as lower triangle: + // row 0: [1] (col 0) + // row 1: [2, 3] (cols 0, 1) + // row 2: [4, 5] (cols 1, 2) + mjtNum mat[] = {1, 2, 3, 4, 5}; + int rownnz[] = {1, 2, 2}; + int rowadr[] = {0, 1, 3}; + int colind[] = {0, 0, 1, 1, 2}; + + mjtNum dense[9]; + mju_sym2dense(dense, mat, 3, rownnz, rowadr, colind); + + mjtNum expected[] = {1, 2, 0, 2, 3, 4, 0, 4, 5}; + EXPECT_EQ(AsVector(dense, 9), AsVector(expected, 9)); +} + +TEST_F(EngineUtilSparseTest, MjuSym2DenseWithUpper) { + mjtNum mat[] = {1, 999, 2, 3, 4, 5}; + int rownnz[] = {2, 2, 2}; + int rowadr[] = {0, 2, 4}; + int colind[] = {0, 1, 0, 1, 1, 2}; + + mjtNum dense[9]; + mju_sym2dense(dense, mat, 3, rownnz, rowadr, colind); + + mjtNum expected[] = {1, 2, 0, + 2, 3, 4, + 0, 4, 5}; + EXPECT_EQ(AsVector(dense, 9), AsVector(expected, 9)); +} + // helper: run split-col approach and return dense result static void SqrMatTDSplitCol( std::vector& dense_result, int nr, int nc, diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 211665e7..c7cbbe66 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -7514,6 +7514,9 @@ public static unsafe extern int mju_dense2sparse(double* res, double* mat, int n [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mju_sparse2dense(double* res, double* mat, int nr, int nc, int* rownnz, int* rowadr, int* colind); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +public static unsafe extern void mju_sym2dense(double* res, double* mat, int n, int* rownnz, int* rowadr, int* colind); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mju_rotVecQuat(double* res, double* vec, double* quat); diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 2dd89505..eddd2024 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -10638,6 +10638,15 @@ mjtNum mju_sum_wrapper(const NumberArray& vec, int n) { return mju_sum(vec_.data(), n); } +void mju_sym2dense_wrapper(const val& res, const NumberArray& mat, int n, const NumberArray& rownnz, const NumberArray& rowadr, const NumberArray& colind) { + UNPACK_VALUE(mjtNum, res); + UNPACK_ARRAY(mjtNum, mat); + UNPACK_ARRAY(int, rownnz); + UNPACK_ARRAY(int, rowadr); + UNPACK_ARRAY(int, colind); + mju_sym2dense(res_.data(), mat_.data(), n, rownnz_.data(), rowadr_.data(), colind_.data()); +} + void mju_symmetrize_wrapper(const val& res, const NumberArray& mat, int n) { UNPACK_VALUE(mjtNum, res); UNPACK_ARRAY(mjtNum, mat); @@ -13461,6 +13470,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { function("mju_subFrom3", &mju_subFrom3_wrapper); function("mju_subQuat", &mju_subQuat_wrapper); function("mju_sum", &mju_sum_wrapper); + function("mju_sym2dense", &mju_sym2dense_wrapper); function("mju_symmetrize", &mju_symmetrize_wrapper); function("mju_transformSpatial", &mju_transformSpatial_wrapper); function("mju_transpose", &mju_transpose_wrapper);