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_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);