From e56b31e98ff07c84e630b55bd285860347439d86 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Wed, 14 Jan 2026 11:17:47 -0800 Subject: [PATCH] Reduce flexedge_J size from nflexedge x nv to the effective sparse size. PiperOrigin-RevId: 856290141 Change-Id: Ide297d1b6f0e3aa07e3e20bf9ab44b6501097695 --- doc/includes/references.h | 5 +- include/mujoco/mjdata.h | 2 +- include/mujoco/mjmodel.h | 3 +- include/mujoco/mjxmacro.h | 5 +- mjx/mujoco/mjx/_src/io.py | 2 +- python/mujoco/introspect/structs.py | 9 ++- src/engine/engine_core_constraint.c | 24 ++++---- src/engine/engine_core_smooth.c | 65 ++++++++-------------- src/engine/engine_core_util.c | 3 +- src/engine/engine_core_util.h | 2 +- src/engine/engine_derivative.c | 10 +--- src/engine/engine_forward.c | 10 +--- src/engine/engine_io.c | 11 ++-- src/engine/engine_io.h | 2 +- src/engine/engine_passive.c | 19 +++---- src/engine/engine_print.c | 14 ++--- src/engine/engine_setconst.c | 12 ++-- src/user/user_model.cc | 27 ++++++++- src/user/user_model.h | 1 + test/engine/engine_core_constraint_test.cc | 2 +- unity/Runtime/Bindings/MjBindings.cs | 1 + wasm/codegen/generated/bindings.cc | 11 +++- 22 files changed, 123 insertions(+), 117 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 98bed114..4251e296 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -272,7 +272,7 @@ struct mjData_ { // computed by mj_fwdPosition/mj_flex mjtNum* flexvert_xpos; // Cartesian flex vertex positions (nflexvert x 3) mjtNum* flexelem_aabb; // flex element bounding boxes (center, size) (nflexelem x 6) - mjtNum* flexedge_J; // flex edge Jacobian (nflexedge x nv) + mjtNum* flexedge_J; // flex edge Jacobian (nJfe x 1) mjtNum* flexedge_length; // flex edge lengths (nflexedge x 1) mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6) @@ -1033,6 +1033,7 @@ struct mjModel_ { int nflexshelldata; // number of shell fragment vertex ids in all flexes int nflexevpair; // number of element-vertex pairs in all flexes int nflextexcoord; // number of vertices with texture coordinates + int nJfe; // number of non-zeros in sparse flexedge Jacobian matrix int nmesh; // number of meshes int nmeshvert; // number of vertices in all meshes int nmeshnormal; // number of normals in all meshes @@ -1341,7 +1342,7 @@ struct mjModel_ { int* flex_bvhnum; // number of bounding volumes (nflex x 1) int* flexedge_J_rownnz; // number of non-zeros in Jacobian row (nflexedge x 1) int* flexedge_J_rowadr; // row start address in colind array (nflexedge x 1) - int* flexedge_J_colind; // column indices in sparse Jacobian (nflexedge x nv) + int* flexedge_J_colind; // column indices in sparse Jacobian (nJfe x 1) float* flex_rgba; // rgba when material is omitted (nflex x 4) float* flex_texcoord; // vertex texture coordinates (nflextexcoord x 2) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 57337d69..202823cc 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -306,7 +306,7 @@ struct mjData_ { // computed by mj_fwdPosition/mj_flex mjtNum* flexvert_xpos; // Cartesian flex vertex positions (nflexvert x 3) mjtNum* flexelem_aabb; // flex element bounding boxes (center, size) (nflexelem x 6) - mjtNum* flexedge_J; // flex edge Jacobian (nflexedge x nv) + mjtNum* flexedge_J; // flex edge Jacobian (nJfe x 1) mjtNum* flexedge_length; // flex edge lengths (nflexedge x 1) mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 45e99516..c951de0b 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -701,6 +701,7 @@ struct mjModel_ { int nflexshelldata; // number of shell fragment vertex ids in all flexes int nflexevpair; // number of element-vertex pairs in all flexes int nflextexcoord; // number of vertices with texture coordinates + int nJfe; // number of non-zeros in sparse flexedge Jacobian matrix int nmesh; // number of meshes int nmeshvert; // number of vertices in all meshes int nmeshnormal; // number of normals in all meshes @@ -1009,7 +1010,7 @@ struct mjModel_ { int* flex_bvhnum; // number of bounding volumes (nflex x 1) int* flexedge_J_rownnz; // number of non-zeros in Jacobian row (nflexedge x 1) int* flexedge_J_rowadr; // row start address in colind array (nflexedge x 1) - int* flexedge_J_colind; // column indices in sparse Jacobian (nflexedge x nv) + int* flexedge_J_colind; // column indices in sparse Jacobian (nJfe x 1) float* flex_rgba; // rgba when material is omitted (nflex x 4) float* flex_texcoord; // vertex texture coordinates (nflextexcoord x 2) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index c4db5515..94f8ad01 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -96,6 +96,7 @@ X( nflexshelldata ) \ X( nflexevpair ) \ X( nflextexcoord ) \ + X( nJfe ) \ X( nmesh ) \ X( nmeshvert ) \ X( nmeshnormal ) \ @@ -399,7 +400,7 @@ X ( int, flex_bvhnum, nflex, 1 ) \ X ( int, flexedge_J_rownnz, nflexedge, 1 ) \ X ( int, flexedge_J_rowadr, nflexedge, 1 ) \ - X ( int, flexedge_J_colind, nflexedge, MJ_M(nv) ) \ + X ( int, flexedge_J_colind, nJfe, 1 ) \ X ( float, flex_rgba, nflex, 4 ) \ X ( float, flex_texcoord, nflextexcoord, 2 ) @@ -733,7 +734,7 @@ X ( mjtNum, cinert, nbody, 10 ) \ X ( mjtNum, flexvert_xpos, nflexvert, 3 ) \ X ( mjtNum, flexelem_aabb, nflexelem, 6 ) \ - X ( mjtNum, flexedge_J, nflexedge, MJ_M(nv) ) \ + X ( mjtNum, flexedge_J, nJfe, 1 ) \ X ( mjtNum, flexedge_length, nflexedge, 1 ) \ X ( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ X ( int, ten_wrapadr, ntendon, 1 ) \ diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 68d1a69a..29cf5d7d 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -754,7 +754,7 @@ def _make_data_c( 'light_xdir': (m.nlight, 3, float_), 'flexvert_xpos': (nflexvert, 3, float_), 'flexelem_aabb': (nflexelem, 6, float_), - 'flexedge_J': (nflexedge, m.nv, float_), + 'flexedge_J': (m.nJfe, float_), 'flexedge_length': (nflexedge, float_), 'ten_J_rownnz': (m.ntendon, np.int32), 'ten_J_rowadr': (m.ntendon, np.int32), diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index 808aed09..f518de53 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -992,6 +992,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of vertices with texture coordinates', ), + StructFieldDecl( + name='nJfe', + type=ValueType(name='int'), + doc='number of non-zeros in sparse flexedge Jacobian matrix', + ), StructFieldDecl( name='nmesh', type=ValueType(name='int'), @@ -2993,7 +2998,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='int'), ), doc='column indices in sparse Jacobian', - array_extent=('nflexedge', 'nv'), + array_extent=('nJfe',), ), StructFieldDecl( name='flex_rgba', @@ -5594,7 +5599,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='mjtNum'), ), doc='flex edge Jacobian', - array_extent=('nflexedge', 'nv'), + array_extent=('nJfe',), ), StructFieldDecl( name='flexedge_length', diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index bc7ba7f3..c90f35cc 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -447,7 +447,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // compute Jacobian difference (opposite of contact: 0 - 1) NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0], - jac[1], jac[0], jacdif, NULL, NULL, NULL); + jac[1], jac[0], jacdif, NULL, NULL, NULL, issparse); // copy difference into jac[0] mju_copy(jac[0], jacdif, 3*NV); @@ -483,7 +483,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // compute error Jacobian (opposite of contact: 0 - 1) NV = mj_jacDifPair(m, d, chain, body_id[1], body_id[0], pos[1], pos[0], jac[1], jac[0], jacdif, - jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv); + jac[1]+3*nv, jac[0]+3*nv, jacdif+3*nv, issparse); // copy difference into jac[0], compress translation:rotation if sparse mju_copy(jac[0], jacdif, 3*NV); @@ -628,13 +628,17 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) { // add constraint: sparse or dense if (issparse) { mj_addConstraint(m, d, d->flexedge_J+m->flexedge_J_rowadr[e], cpos, 0, 0, - 1, mjCNSTR_EQUALITY, i, - m->flexedge_J_rownnz[e], - m->flexedge_J_colind+m->flexedge_J_rowadr[e]); + 1, mjCNSTR_EQUALITY, i, + m->flexedge_J_rownnz[e], + m->flexedge_J_colind+m->flexedge_J_rowadr[e]); } else { - mj_addConstraint(m, d, d->flexedge_J+e*nv, cpos, 0, 0, - 1, mjCNSTR_EQUALITY, i, - 0, NULL); + mju_zero(jac[0], nv); // reuse first row of jac[0] + int rowadr = m->flexedge_J_rowadr[e]; + int rownnz = m->flexedge_J_rownnz[e]; + for (int k=0; kflexedge_J_colind[rowadr+k]] = d->flexedge_J[rowadr+k]; + } + mj_addConstraint(m, d, jac[0], cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL); } } break; @@ -891,10 +895,10 @@ int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int di // compute Jacobian differences if (dim > 3) { return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos, - jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr); + jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr, mj_isSparse(m)); } else { return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos, - jac1p, jac2p, jacdifp, NULL, NULL, NULL); + jac1p, jac2p, jacdifp, NULL, NULL, NULL, mj_isSparse(m)); } } diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 3b206438..2926bb73 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -535,9 +535,8 @@ void mj_updateDynamicBVH(const mjModel* m, mjData* d, int bvhadr, int bvhnum) { // compute flex-related quantities void mj_flex(const mjModel* m, mjData* d) { - int nv = m->nv, issparse = mj_isSparse(m); + int nv = m->nv; int* rowadr = m->flexedge_J_rowadr, *rownnz = m->flexedge_J_rownnz; - mjtNum* J = d->flexedge_J; // skip if no flexes if (!m->nflex) { @@ -655,15 +654,11 @@ void mj_flex(const mjModel* m, mjData* d) { mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum); - int* chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL; + int* chain = mjSTACKALLOC(d, nv, int); - // clear Jacobian: sparse or dense - if (issparse) { - mju_zeroInt(rowadr, m->nflexedge); - mju_zeroInt(rownnz, m->nflexedge); - } else { - mju_zero(J, m->nflexedge*nv); - } + // clear Jacobian + mju_zeroInt(rowadr, m->nflexedge); + mju_zeroInt(rownnz, m->nflexedge); // compute lengths and Jacobians of edges for (int f=0; f < m->nflex; f++) { @@ -699,40 +694,26 @@ void mj_flex(const mjModel* m, mjData* d) { continue; } - // sparse edge Jacobian - if (issparse) { - // set rowadr - if (ebase+e > 0) { - rowadr[ebase+e] = rowadr[ebase+e-1] + rownnz[ebase+e-1]; - } - - // get endpoint Jacobians, subtract - int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2, - jac1, jac2, jacdif, NULL, NULL, NULL); - - // no dofs: skip - if (!NV) { - continue; - } - - // apply chain rule to compute edge Jacobian - mju_mulMatTVec(J + rowadr[ebase+e], jacdif, vec, 3, NV); - - // copy sparsity info - rownnz[ebase+e] = NV; - mju_copyInt(m->flexedge_J_colind + rowadr[ebase+e], chain, NV); + // set rowadr + if (ebase+e > 0) { + rowadr[ebase+e] = rowadr[ebase+e-1] + rownnz[ebase+e-1]; } - // dense edge Jacobian - else { - // get endpoint Jacobians, subtract - mj_jac(m, d, jac1, NULL, pos1, b1); - mj_jac(m, d, jac2, NULL, pos2, b2); - mju_sub(jacdif, jac2, jac1, 3*nv); + // get endpoint Jacobians, subtract + int NV = mj_jacDifPair(m, d, chain, b1, b2, pos1, pos2, + jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); - // apply chain rule to compute edge Jacobian - mju_mulMatTVec(J + (ebase+e)*nv, jacdif, vec, 3, nv); + // no dofs: skip + if (!NV) { + continue; } + + // apply chain rule to compute edge Jacobian + mju_mulMatTVec(d->flexedge_J + rowadr[ebase+e], jacdif, vec, 3, NV); + + // copy sparsity info + rownnz[ebase+e] = NV; + mju_copyInt(m->flexedge_J_colind + rowadr[ebase+e], chain, NV); } } @@ -903,7 +884,7 @@ void mj_tendon(const mjModel* m, mjData* d) { // get endpoint Jacobians, subtract int NV = mj_jacDifPair(m, d, chain, wbody[k], wbody[k+1], wpnt+3*k, wpnt+3*k+3, - jac1, jac2, jacdif, NULL, NULL, NULL); + jac1, jac2, jacdif, NULL, NULL, NULL, /*issparse=*/1); // no dofs: skip if (!NV) { @@ -1526,7 +1507,7 @@ void mj_transmission(const mjModel* m, mjData* d) { // get Jacobian difference int NV = mj_jacDifPair(m, d, chain, b1, b2, con->pos, con->pos, - jac1p, jac2p, jacdifp, NULL, NULL, NULL); + jac1p, jac2p, jacdifp, NULL, NULL, NULL, issparse); // project Jacobian along the normal of the contact frame mju_mulMatMat(jac, con->frame, jacdifp, 1, 3, NV); diff --git a/src/engine/engine_core_util.c b/src/engine/engine_core_util.c index cfe3d099..da75ec8d 100644 --- a/src/engine/engine_core_util.c +++ b/src/engine/engine_core_util.c @@ -438,9 +438,8 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d, int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain, int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3], mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp, - mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr) { + mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse) { int issimple = (m->body_simple[b1] && m->body_simple[b2]); - int issparse = mj_isSparse(m); int NV = m->nv; // skip if no DOFs diff --git a/src/engine/engine_core_util.h b/src/engine/engine_core_util.h index 3b4d7465..d44f175f 100644 --- a/src/engine/engine_core_util.h +++ b/src/engine/engine_core_util.h @@ -89,7 +89,7 @@ void mj_jacSparseSimple(const mjModel* m, const mjData* d, MJAPI int mj_jacDifPair(const mjModel* m, const mjData* d, int* chain, int b1, int b2, const mjtNum pos1[3], const mjtNum pos2[3], mjtNum* jac1p, mjtNum* jac2p, mjtNum* jacdifp, - mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr); + mjtNum* jac1r, mjtNum* jac2r, mjtNum* jacdifr, int issparse); // dense or sparse weighted sum of multiple body Jacobians at same point int mj_jacSum(const mjModel* m, mjData* d, int* chain, diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index e810601a..7a11c1f7 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -1495,13 +1495,9 @@ void mjd_passive_vel(const mjModel* m, mjData* d) { continue; } - // add sparse or dense - if (mj_isSparse(m)) { - addJTBJSparse(m, d, d->flexedge_J, &B, 1, e, - m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind); - } else { - addJTBJ(m, d, d->flexedge_J+e*nv, &B, 1); - } + // always sparse + addJTBJSparse(m, d, d->flexedge_J, &B, 1, e, + m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind); } } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index b62babf7..03e86c61 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -218,13 +218,9 @@ void mj_fwdPosition(const mjModel* m, mjData* d) { void mj_fwdVelocity(const mjModel* m, mjData* d) { TM_START; - // flexedge velocity: dense or sparse - if (mj_isSparse(m)) { - mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, - m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL); - } else { - mju_mulMatVec(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, m->nv); - } + // flexedge velocity: always sparse + mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, + m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL); // tendon velocity: dense or sparse if (mj_isSparse(m)) { diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 538e0afd..902361b5 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -218,9 +218,9 @@ static void freeModelBuffers(mjModel* m) { void mj_makeModel(mjModel** dest, int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic, int noct, int njnt, int ntree, - int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, - int nlight, int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem, - int nflexelemdata, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, + int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight, + int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata, + int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, int nJfe, int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface, int nmeshgraph, int nmeshpoly, int nmeshpolyvert, int nmeshpolymap, int nskin, int nskinvert, int nskintexvert, int nskinface, @@ -278,6 +278,7 @@ void mj_makeModel(mjModel** dest, m->nflexshelldata = nflexshelldata; m->nflexevpair = nflexevpair; m->nflextexcoord = nflextexcoord; + m->nJfe = nJfe; m->nmesh = nmesh; m->nmeshvert = nmeshvert; m->nmeshnormal = nmeshnormal; @@ -406,7 +407,7 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) { src->nM, src->nB, src->nC, src->nD, src->ngeom, src->nsite, src->ncam, src->nlight, src->nflex, src->nflexnode, src->nflexvert, src->nflexedge, src->nflexelem, src->nflexelemdata, src->nflexelemedge, - src->nflexshelldata, src->nflexevpair, src->nflextexcoord, src->nmesh, + src->nflexshelldata, src->nflexevpair, src->nflextexcoord, src->nJfe, src->nmesh, src->nmeshvert, src->nmeshnormal, src->nmeshtexcoord, src->nmeshface, src->nmeshgraph, src->nmeshpoly, src->nmeshpolyvert, src->nmeshpolymap, src->nskin, src->nskinvert, src->nskintexvert, src->nskinface, @@ -597,7 +598,7 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) { ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55], ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62], ints[63], ints[64], ints[65], ints[66], ints[67], ints[68], ints[69], - ints[70], ints[71], ints[72], ints[73], ints[74]); + ints[70], ints[71], ints[72], ints[73], ints[74], ints[75]); // read mjModel mjtSize fields mjtSize sizes[8]; diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index a5410f86..e4a8e355 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -51,7 +51,7 @@ void mj_makeModel(mjModel** dest, int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic, int noct, int njnt, int ntree, int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata, - int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, int nmesh, + int nJfe, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface, int nmeshgraph, int nmeshpoly, int nmeshpolyvert, int nmeshpolymap, int nskin, int nskinvert, int nskintexvert, int nskinface, int nskinbone, int nskinbonevert, int nhfield, int nhfielddata, int ntex, int ntexdata, diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 535a99f9..3ed27063 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -412,18 +412,13 @@ static void mj_springdamper(const mjModel* m, mjData* d) { mjtNum frc_spring = stiffness * (m->flexedge_length0[e] - d->flexedge_length[e]); mjtNum frc_damper = -damping * d->flexedge_velocity[e]; - // transform to joint torque, add to qfrc_{spring, damper}: dense or sparse - if (issparse) { - int end = m->flexedge_J_rowadr[e] + m->flexedge_J_rownnz[e]; - for (int j=m->flexedge_J_rowadr[e]; j < end; j++) { - int colind = m->flexedge_J_colind[j]; - mjtNum J = d->flexedge_J[j]; - d->qfrc_spring[colind] += J * frc_spring; - d->qfrc_damper[colind] += J * frc_damper; - } - } else { - if (frc_spring) mju_addToScl(d->qfrc_spring, d->flexedge_J+e*nv, frc_spring, nv); - if (frc_damper) mju_addToScl(d->qfrc_damper, d->flexedge_J+e*nv, frc_damper, nv); + // transform to joint torque, add to qfrc_{spring, damper}: always sparse + int end = m->flexedge_J_rowadr[e] + m->flexedge_J_rownnz[e]; + for (int j=m->flexedge_J_rowadr[e]; j < end; j++) { + int colind = m->flexedge_J_colind[j]; + mjtNum J = d->flexedge_J[j]; + d->qfrc_spring[colind] += J * frc_spring; + d->qfrc_damper[colind] += J * frc_damper; } } } diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 78071b01..90c04b06 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1280,15 +1280,11 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printArray2d("FLEXVERT_XPOS", m->nflexvert, 3, d->flexvert_xpos, fp, float_format); printArray2d("FLEXELEM_AABB", m->nflexelem, 6, d->flexelem_aabb, fp, float_format); - if (!mj_isSparse(m)) { - printArray2d("FLEXEDGE_J", m->nflexedge, m->nv, d->flexedge_J, fp, float_format); - } else { - mj_printSparsity("FLEXEDGE_J: flex edge connectivity", m->nflexedge, m->nv, - m->flexedge_J_rowadr, NULL, m->flexedge_J_rownnz, NULL, m->flexedge_J_colind, - fp); - printSparse("FLEXEDGE_J", d->flexedge_J, m->nflexedge, m->flexedge_J_rownnz, - m->flexedge_J_rowadr, m->flexedge_J_colind, fp, float_format); - } + mj_printSparsity("FLEXEDGE_J: flex edge connectivity", m->nflexedge, m->nv, + m->flexedge_J_rowadr, NULL, m->flexedge_J_rownnz, NULL, m->flexedge_J_colind, + fp); + printSparse("FLEXEDGE_J", d->flexedge_J, m->nflexedge, m->flexedge_J_rownnz, + m->flexedge_J_rowadr, m->flexedge_J_colind, fp, float_format); printArray2d("FLEXEDGE_LENGTH", m->nflexedge, 1, d->flexedge_length, fp, float_format); printArray2d("TEN_LENGTH", m->ntendon, 1, d->ten_length, fp, float_format); diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 60a49135..56200cc6 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -468,14 +468,10 @@ static void set0(mjModel* m, mjData* d) { // handle general edge else { // make dense vector into tmp - if (mj_isSparse(m)) { - mju_zero(tmp, nv); - int end = m->flexedge_J_rowadr[i] + m->flexedge_J_rownnz[i]; - for (int j=m->flexedge_J_rowadr[i]; j < end; j++) { - tmp[m->flexedge_J_colind[j]] = d->flexedge_J[j]; - } - } else { - mju_copy(tmp, d->flexedge_J+i*nv, nv); + mju_zero(tmp, nv); + int end = m->flexedge_J_rowadr[i] + m->flexedge_J_rownnz[i]; + for (int j=m->flexedge_J_rowadr[i]; j < end; j++) { + tmp[m->flexedge_J_colind[j]] = d->flexedge_J[j]; } // solve into tmp+nv diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 624b77c2..3bba27bf 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -34,6 +34,7 @@ #include #include #include +#include #include #include @@ -1153,6 +1154,7 @@ void mjCModel::Clear() { nflexshelldata = 0; nflexevpair = 0; nflextexcoord = 0; + nJfe = 0; nmeshvert = 0; nmeshnormal = 0; nmeshtexcoord = 0; @@ -2141,6 +2143,29 @@ void mjCModel::SetSizes() { nflexshelldata += (int)flexes_[i]->shell.size(); nflexevpair += (int)flexes_[i]->evpair.size()/2; nflextexcoord += (flexes_[i]->HasTexcoord() ? flexes_[i]->get_texcoord().size()/2 : 0); + if (flexes_[i]->interpolated || flexes_[i]->rigid) { + continue; + } + + // count number of non-zero elements in the edge Jacobian matrix + for (const auto& edge : flexes_[i]->edge) { + mjCBody* b1 = bodies_[flexes_[i]->vertbodyid[edge.first]]; + mjCBody* b2 = bodies_[flexes_[i]->vertbodyid[edge.second]]; + std::unordered_set bodies_in_jac; + while (b1 || b2) { + if (b1) { + bodies_in_jac.insert(b1); + b1 = b1->parent; + } + if (b2) { + bodies_in_jac.insert(b2); + b2 = b2->parent; + } + } + for (mjCBody* b : bodies_in_jac) { + nJfe += b->dofnum; + } + } } // mesh counts @@ -4831,7 +4856,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { mj_makeModel(&m, nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, noct, njnt, ntree, nM, nB, nC, nD, ngeom, nsite, ncam, nlight, nflex, nflexnode, nflexvert, nflexedge, nflexelem, - nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord, + nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord, nJfe, nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph, nmeshpoly, nmeshpolyvert, nmeshpolymap, nskin, nskinvert, nskintexvert, nskinface, nskinbone, nskinbonevert, nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude, diff --git a/src/user/user_model.h b/src/user/user_model.h index b0ab37c8..a4af02f5 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -99,6 +99,7 @@ class mjCModel_ : public mjsElement { int nflexshelldata; // number of shell fragment vertex ids in all flexes int nflexevpair; // number of element-vertex pairs in all flexes int nflextexcoord; // number of vertex texture coordinates in all flexes + int nJfe; // number of non-zeros in sparse flex constraint Jacobian int nmeshvert; // number of vertices in all meshes int nmeshnormal; // number of normals in all meshes int nmeshtexcoord; // number of texture coordinates in all meshes diff --git a/test/engine/engine_core_constraint_test.cc b/test/engine/engine_core_constraint_test.cc index db7227b5..d37a4f15 100644 --- a/test/engine/engine_core_constraint_test.cc +++ b/test/engine/engine_core_constraint_test.cc @@ -133,7 +133,7 @@ TEST_F(CoreConstraintTest, WeldRotJacobian) { // rotational Jacobian difference mj_jacDifPair(model, data, NULL, 2, 1, point, point, - NULL, NULL, NULL, jac0, jac1, jacdif); + NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model)); // formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1 mjtNum axis[3], quat3[4], quat4[4]; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index fe14ba1e..dc0e9e1f 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5313,6 +5313,7 @@ public unsafe struct mjModel_ { public int nflexshelldata; public int nflexevpair; public int nflextexcoord; + public int nJfe; public int nmesh; public int nmeshvert; public int nmeshnormal; diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index def6244a..cd6231dc 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -3579,6 +3579,12 @@ struct MjModel { void set_nflextexcoord(int value) { ptr_->nflextexcoord = value; } + int nJfe() const { + return ptr_->nJfe; + } + void set_nJfe(int value) { + ptr_->nJfe = value; + } int nmesh() const { return ptr_->nmesh; } @@ -4558,7 +4564,7 @@ struct MjModel { return emscripten::val(emscripten::typed_memory_view(ptr_->nflexedge, ptr_->flexedge_J_rowadr)); } emscripten::val flexedge_J_colind() const { - return emscripten::val(emscripten::typed_memory_view(ptr_->nflexedge * ptr_->nv, ptr_->flexedge_J_colind)); + return emscripten::val(emscripten::typed_memory_view(ptr_->nJfe, ptr_->flexedge_J_colind)); } emscripten::val flex_rgba() const { return emscripten::val(emscripten::typed_memory_view(ptr_->nflex * 4, ptr_->flex_rgba)); @@ -6295,7 +6301,7 @@ struct MjData { return emscripten::val(emscripten::typed_memory_view(model->nflexelem * 6, ptr_->flexelem_aabb)); } emscripten::val flexedge_J() const { - return emscripten::val(emscripten::typed_memory_view(model->nflexedge * model->nv, ptr_->flexedge_J)); + return emscripten::val(emscripten::typed_memory_view(model->nJfe, ptr_->flexedge_J)); } emscripten::val flexedge_length() const { return emscripten::val(emscripten::typed_memory_view(model->nflexedge, ptr_->flexedge_length)); @@ -11470,6 +11476,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("nB", &MjModel::nB, &MjModel::set_nB, reference()) .property("nC", &MjModel::nC, &MjModel::set_nC, reference()) .property("nD", &MjModel::nD, &MjModel::set_nD, reference()) + .property("nJfe", &MjModel::nJfe, &MjModel::set_nJfe, reference()) .property("nJmom", &MjModel::nJmom, &MjModel::set_nJmom, reference()) .property("nM", &MjModel::nM, &MjModel::set_nM, reference()) .property("na", &MjModel::na, &MjModel::set_na, reference())