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

Reduce flex stiffness and bending memory size to the exact amount needed.

Browse Source 
PiperOrigin-RevId: 911774147
Change-Id: I3b18eba23e75ec91ece8342945f25e0cd7c63a1a
This commit is contained in:
Alessio Quaglino
2026-05-06 23:52:46 -07:00
committed by Copybara-Service
parent b6aac76cc4
commit 0c05215e18
5 changed files with 70 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+14 -6
View File
@@ -832,9 +832,14 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
// read eigenmode data from flex_stiffness
int ndof_elem = 3 * npe;
const mjtNum* k_elem = m->flex_stiffness + m->flex_stiffnessadr[f]
+ elem_idx * ndof_elem * ndof_elem;
int neig = (int)k_elem[0];
int stiffnessadr = m->flex_stiffnessadr[f];
int neig = 0;
const mjtNum* k_elem = NULL;
if (stiffnessadr >= 0) {
k_elem = m->flex_stiffness + stiffnessadr
+ elem_idx * ndof_elem * ndof_elem;
neig = (int)k_elem[0];
}
// compute displacement in corotational frame
mjtNum* displ_e = mjSTACKALLOC(d, ndof_elem, mjtNum);
@@ -2387,9 +2392,12 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
// read eigenmode count from flex_stiffness
int ndof_elem = 3 * npe;
const mjtNum* k_elem = m->flex_stiffness + m->flex_stiffnessadr[f]
+ elem_idx * ndof_elem * ndof_elem;
size = (int)k_elem[0]; // neig stored as first element
size = 0;
if (m->flex_stiffnessadr[f] >= 0) {
const mjtNum* k_elem = m->flex_stiffness + m->flex_stiffnessadr[f]
+ elem_idx * ndof_elem * ndof_elem;
size = (int)k_elem[0]; // neig stored as first element
}
if (nnz) {
// get element node body IDs
src/engine/engine_derivative.c
+5 -1
View File
@@ -951,7 +951,11 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
}
// get stiffness and damping
mjtNum* K = m->flex_stiffness + m->flex_stiffnessadr[f];
int stiffnessadr = m->flex_stiffnessadr[f];
if (stiffnessadr < 0) {
continue;
}
mjtNum* K = m->flex_stiffness + stiffnessadr;
// skip if rigid or no stiffness
if (m->flex_rigid[f] || K[0] == 0) {
src/engine/engine_passive.c
+28 -7
View File
@@ -62,7 +62,12 @@ static void inline GradSquaredLengths(mjtNum gradient[6][2][3],
// passive forces for interpolated flex (stretch + bending)
static void mj_flexPassiveInterp(const mjModel* m, mjData* d, int f,
int enbl_spring, int enbl_damper) {
mjtNum* k = m->flex_stiffness + m->flex_stiffnessadr[f];
int stiffnessadr = m->flex_stiffnessadr[f];
if (stiffnessadr < 0) {
return;
}
mjtNum* k = m->flex_stiffness + stiffnessadr;
int nodenum = m->flex_nodenum[f];
int order = m->flex_interp[f];
@@ -230,8 +235,17 @@ static inline mjtNum mju_dphi2D(mjtNum s0, int l0, mjtNum s1, int l1,
// accuracy for elements with varying curvature.
static void mj_flexPassiveBendInterp(const mjModel* m, mjData* d, int f,
int enbl_spring, int enbl_damper) {
if (m->flex_interp[f] >= 0) {
return;
}
int bendingadr = m->flex_bendingadr[f];
if (bendingadr < 0) {
return;
}
// read bending edge data
const mjtNum* bdata = m->flex_bending + m->flex_bendingadr[f];
const mjtNum* bdata = m->flex_bending + bendingadr;
int nedge = (int)bdata[0];
if (nedge == 0) return;
@@ -398,10 +412,15 @@ static void mj_flexPassiveBend(const mjModel* m, mjData* d, int f,
return;
}
int bendingadr = m->flex_bendingadr[f];
if (bendingadr < 0) {
return;
}
int edgenum = m->flex_edgenum[f];
mjtNum* xpos = d->flexvert_xpos + 3*m->flex_vertadr[f];
int* bodyid = m->flex_vertbodyid + m->flex_vertadr[f];
mjtNum* b = m->flex_bending + 17*m->flex_edgeadr[f];
mjtNum* b = m->flex_bending + bendingadr;
for (int e = 0; e < edgenum; e++) {
const int* edge = m->flex_edge + 2*(e+m->flex_edgeadr[f]);
@@ -469,7 +488,11 @@ static void mj_flexPassiveBend(const mjModel* m, mjData* d, int f,
// passive forces for flex stretch
static void mj_flexPassiveStretch(const mjModel* m, mjData* d, int f,
int enbl_spring, int enbl_damper) {
mjtNum* k = m->flex_stiffness + m->flex_stiffnessadr[f];
int stiffnessadr = m->flex_stiffnessadr[f];
if (stiffnessadr < 0) {
return;
}
mjtNum* k = m->flex_stiffness + stiffnessadr;
if (k[0] == 0) {
return;
}
@@ -660,9 +683,7 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
mj_flexPassiveInterp(m, d, f, enbl_spring, enbl_damper);
// interpolated shell bending forces
if (m->flex_interp[f] < 0) {
mj_flexPassiveBendInterp(m, d, f, enbl_spring, enbl_damper);
}
mj_flexPassiveBendInterp(m, d, f, enbl_spring, enbl_damper);
} else {
// add bending forces
mj_flexPassiveBend(m, d, f, enbl_spring, enbl_damper);
src/user/user_mesh.cc
+8
View File
@@ -4918,6 +4918,14 @@ void mjCFlex::Compile(const mjVFS* vfs) {
stiffness_cached = LoadCachedStiffness();
}
// check if any strain equality references this flex
for (auto* equality : model->Equalities()) {
if (equality->spec.type == mjEQ_FLEXSTRAIN && *equality->spec.name1 == name) {
has_strain_eq = true;
break;
}
}
if (!stiffness_cached && interpolated && (young > 0 || has_strain_eq)) {
// use young=1 for strain constraints (eigenvectors are geometry-only)
double K_young = has_strain_eq ? 1e1 : young;
src/user/user_model.cc
+15 -67
View File
@@ -2106,21 +2106,6 @@ static size_t getpathslength(std::vector<T> list) {
return result;
}
// compute extra stiffness/bending array size for an interpolated flex
static int flexInterpExtraSize(int order, const int cellcount[3], bool shell) {
int npe, nelem;
int cx = cellcount[0], cy = cellcount[1], cz = cellcount[2];
if (shell) {
npe = (int)pow(order + 1, 2);
nelem = 2*(cy*cz + cx*cz + cx*cy);
} else {
npe = (int)pow(order + 1, 3);
nelem = cx * cy * cz;
}
int ndof_elem = 3 * npe;
return nelem * ndof_elem * ndof_elem;
}
// set array sizes
void mjCModel::SetSizes() {
// set from object list sizes
@@ -2194,8 +2179,6 @@ void mjCModel::SetSizes() {
}
nbvh = nbvhstatic + nbvhdynamic;
int extra_stiffness_size = 0;
int extra_bending_size = 0;
// flex counts
for (int i=0; i < nflex; i++) {
nflexnode += flexes_[i]->nnode;
@@ -2207,13 +2190,8 @@ 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]->spec.order != 0) {
int extra_size = flexInterpExtraSize(
flexes_[i]->spec.order, flexes_[i]->spec.cellcount,
flexes_[i]->elastic2d != 0);
extra_stiffness_size += extra_size;
extra_bending_size += extra_size;
}
nflexstiffness += flexes_[i]->stiffness.size();
nflexbending += flexes_[i]->bending.size();
if (flexes_[i]->interpolated || flexes_[i]->rigid) {
continue;
}
@@ -2263,10 +2241,6 @@ void mjCModel::SetSizes() {
}
}
}
// TODO: This can be compacted further when we update mjwarp to not rely on
// 21*elem_adr for non-interpolated flexes and 17*edge_adr for bending.
nflexstiffness = nflexelem * 21 + extra_stiffness_size;
nflexbending = nflexedge * 17 + extra_bending_size;
// mesh counts
for (int i=0; i < nmesh; i++) {
@@ -3351,6 +3325,7 @@ int mjCModel::CountNJten(const mjModel* m) {
// copy objects outside kinematic tree
void mjCModel::CopyObjects(mjModel* m) {
mjtSize adr, bone_adr, vert_adr, node_adr, normal_adr, face_adr, texcoord_adr, oct_adr;
mjtSize stiffness_adr, bending_adr;
mjtSize edge_adr, elem_adr, elemdata_adr, elemedge_adr, shelldata_adr, evpair_adr;
mjtSize bonevert_adr, graph_adr, data_adr, bvh_adr;
mjtSize poly_adr, polymap_adr, polyvert_adr;
@@ -3465,10 +3440,8 @@ void mjCModel::CopyObjects(mjModel* m) {
shelldata_adr = 0;
evpair_adr = 0;
texcoord_adr = 0;
int standard_stiffness_size = 21 * m->nflexelem;
int current_extra_stiffness_adr = standard_stiffness_size;
int standard_bending_size = 17 * m->nflexedge;
int current_extra_bending_adr = standard_bending_size;
stiffness_adr = 0;
bending_adr = 0;
for (int i=0; i < nflex; i++) {
// get pointer
mjCFlex* pfl = flexes_[i];
@@ -3491,46 +3464,19 @@ void mjCModel::CopyObjects(mjModel* m) {
mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4);
// elasticity
if (pfl->spec.order == 0) {
m->flex_stiffnessadr[i] = 21 * elem_adr;
if (pfl->stiffness.empty()) {
m->flex_stiffnessadr[i] = -1;
} else {
m->flex_stiffnessadr[i] = current_extra_stiffness_adr;
current_extra_stiffness_adr += flexInterpExtraSize(
pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
}
if (!pfl->stiffness.empty()) {
m->flex_stiffnessadr[i] = stiffness_adr;
mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i],
pfl->stiffness.data(), pfl->stiffness.size());
} else {
int stiff_size;
if (pfl->spec.order == 0) {
stiff_size = 21 * pfl->nelem;
} else {
stiff_size = flexInterpExtraSize(
pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
}
mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], stiff_size);
}
if (pfl->spec.order == 0) {
m->flex_bendingadr[i] = 17 * edge_adr;
if (pfl->bending.empty()) {
m->flex_bendingadr[i] = -1;
} else {
m->flex_bendingadr[i] = current_extra_bending_adr;
current_extra_bending_adr += flexInterpExtraSize(
pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
}
if (!pfl->bending.empty()) {
mjuu_copyvec(m->flex_bending + m->flex_bendingadr[i], pfl->bending.data(), pfl->bending.size());
} else {
int bending_size;
if (pfl->spec.order == 0) {
bending_size = 17 * pfl->nedge;
} else {
bending_size = flexInterpExtraSize(
pfl->spec.order, pfl->spec.cellcount, pfl->elastic2d != 0);
}
mjuu_zerovec(m->flex_bending + m->flex_bendingadr[i], bending_size);
m->flex_bendingadr[i] = bending_adr;
mjuu_copyvec(m->flex_bending + m->flex_bendingadr[i],
pfl->bending.data(), pfl->bending.size());
}
m->flex_damping[i] = (mjtNum)pfl->damping;
@@ -3704,6 +3650,8 @@ void mjCModel::CopyObjects(mjModel* m) {
evpair_adr += (int)pfl->evpair.size()/2;
texcoord_adr += (int)pfl->texcoord_.size()/2;
bvh_adr += pfl->tree.Nbvh();
stiffness_adr += pfl->stiffness.size();
bending_adr += pfl->bending.size();
}
// skins