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

Reduce flexedge_J size from nflexedge x nv to the effective sparse size.

Browse Source 
PiperOrigin-RevId: 856290141
Change-Id: Ide297d1b6f0e3aa07e3e20bf9ab44b6501097695
This commit is contained in:
Alessio Quaglino
2026-01-14 11:17:47 -08:00
committed by Copybara-Service
parent cbc1136502
commit e56b31e98f
22 changed files with 123 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+14 -10
View File
@@ -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; k<rownnz; k++) {
jac[0][m->flexedge_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));
}
}
src/engine/engine_core_smooth.c
+23 -42
View File
@@ -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);
src/engine/engine_core_util.c
+1 -2
View File
@@ -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
src/engine/engine_core_util.h
+1 -1
View File
@@ -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,
src/engine/engine_derivative.c
+3 -7
View File
@@ -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);
}
}
src/engine/engine_forward.c
+3 -7
View File
@@ -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)) {
src/engine/engine_io.c
+6 -5
View File
@@ -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];
src/engine/engine_io.h
+1 -1
View File
@@ -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,
src/engine/engine_passive.c
+7 -12
View File
@@ -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;
}
}
}
src/engine/engine_print.c
+5 -9
View File
@@ -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);
src/engine/engine_setconst.c
+4 -8
View File
@@ -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