chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add mju_sym2dense, document future breakage of mj_fullM

Browse Source 
PiperOrigin-RevId: 910242375
Change-Id: Ibfbdef9cfb66088723499ea257da09aee0d80938
This commit is contained in:
Yuval Tassa
2026-05-04 14:25:23 -07:00
committed by Copybara-Service
parent a51a7bf062
commit 767c607f58
19 changed files with 191 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/APIreference/functions.rst
+9
View File
@@ -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
doc/changelog.rst
+14
View File
@@ -5,6 +5,8 @@ Changelog
Upcoming version (not yet released)
-----------------------------------
General
^^^^^^^
- Added island support for the :ref:`PGS solver<soAlgorithms>`.
- Added support for :ref:`elastic2d<body-flexcomp-elastic2d>` for trilinear and quadratic flex
:ref:`dofs<body-flexcomp-dof>`.
@@ -13,6 +15,18 @@ Upcoming version (not yet released)
(nonzero :ref:`density<option-density>` or :ref:`viscosity<option-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
^^^^^^
doc/includes/references.h
+2
View File
@@ -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]);
include/mujoco/mujoco.h
+4
View File
@@ -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 ---------------------------------------------------
mjx/mujoco/mjx/_src/io.py
+7 -1
View File
@@ -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'])
mjx/mujoco/mjx/_src/io_test.py
+3 -3
View File
@@ -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)
mjx/mujoco/mjx/_src/smooth_test.py
+1 -1
View File
@@ -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')
mjx/mujoco/mjx/_src/support_test.py
+1 -1
View File
@@ -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')
mjx/mujoco/mjx/warp/forward_test.py
+1 -1
View File
@@ -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.
python/mujoco/functions.cc
+31
View File
@@ -1241,6 +1241,37 @@ PYBIND11_MODULE(_functions, pymodule) {
colind.data());
});
DEF_WITH_OMITTED_PY_ARGS(traits::mju_sym2dense, "n")(
pymodule,
[](Eigen::Ref<EigenArrayXX> res,
Eigen::Ref<const EigenVectorX> mat,
Eigen::Ref<const EigenVectorI> rownnz,
Eigen::Ref<const EigenVectorI> rowadr,
Eigen::Ref<const EigenVectorI> 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<traits::mju_rotVecQuat>(pymodule);
Def<traits::mju_negQuat>(pymodule);
python/mujoco/introspect/functions.py
+42
View File
@@ -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',
src/engine/engine_island.c
+1 -1
View File
@@ -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,
src/engine/engine_setconst.c
+1 -1
View File
@@ -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;
}
src/engine/engine_util_sparse.c
+17
View File
@@ -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) {
src/engine/engine_util_sparse.h
+4
View File
@@ -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);
test/engine/engine_support_test.cc
+4 -3
View File
@@ -548,7 +548,8 @@ TEST_F(InertiaTest, mulM) {
// dense M matrix
vector<mjtNum> 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<mjtNum> 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<mjtNum> 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<mjtNum> M_CSR(nv * nv);
test/engine/engine_util_sparse_test.cc
+36
View File
@@ -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<mjtNum>& dense_result, int nr, int nc,
unity/Runtime/Bindings/MjBindings.cs
+3
View File
@@ -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);
wasm/codegen/generated/bindings.cc
+10
View File
@@ -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);