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

Merge pull request #3466 from smallquail:flex-implicit-contact

Browse Source 
PiperOrigin-RevId: 960804180
Change-Id: Ic3eb203ca0d9cf19c9043b46127db342678d94f9
This commit is contained in:
Copybara-Service
2026-08-07 02:36:26 -07:00
parent dca295dde2 2a3554c8a3
commit e144090273
12 changed files with 425 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_constraint.c
+16 -5
View File
@@ -2549,11 +2549,22 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
for (int i=0; i < ncon; i++) {
mjContact* con = d->contact + i;
// skip if passive
if ((con->flex[0] > -1 && m->flex_passive[con->flex[0]]) ||
(con->flex[1] > -1 && m->flex_passive[con->flex[1]])) {
con->efc_address = -1;
con->exclude = 4;
// Passive path: flex-flex (including self-collision) and flex-vs-static-geometry, where every
// dof is a metric-carried flex vertex so the Hessian is assembled in full. Flex-vs-moving-body
// stays on the constraint solver. Passive if either flex asks for it.
{
int f0 = con->flex[0], f1 = con->flex[1];
int wants = (f0 > -1 && m->flex_passive[f0]) || (f1 > -1 && m->flex_passive[f1]);
int ok = (f0 > -1 && f1 > -1); // flex-flex, or a flex with itself
for (int s = 0; s < 2 && !ok; s++) {
if (con->flex[s] < 0 && con->geom[s] > -1) {
ok = (m->body_weldid[m->geom_bodyid[con->geom[s]]] == 0); // welded to the world
}
}
if (wants && ok) {
con->efc_address = -1;
con->exclude = 4;
}
}
// skip if excluded
src/engine/engine_derivative.c
+241 -5
View File
@@ -17,6 +17,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_core_util.h"
#include "engine/engine_crossplatform.h"
@@ -1630,6 +1631,58 @@ mjtBool mjd_flexInterpAssemblable(const mjModel* m) {
// does ANY flex contribute assemblable implicit stiffness? (cheap existence check for the
// solver gate: stretch stiffness on a standard flex, or -- when Krot will be supplied -- an
// operator-processed interp flex)
// does any flex use the passive contact path? Distinct from elasticity: an empty CSR is valid for
// elastic models (matrix-free operators) but means "nothing" for a contact-only flex.
static mjtBool flexPassiveContact_any(const mjModel* m) {
for (int f = 0; f < m->nflex; f++) {
if (!m->flex_interp[f] && !m->flex_rigid[f] && m->flex_dim[f] >= 2 && m->flex_passive[f]) {
return 1;
}
}
return 0;
}
// res += scale * K_contact * vec, where K_contact = sum_c k_c * J_c^T J_c over passive flex
// contacts. Any class that contributes to K must also contribute to the shift -h*K*v (see
// mjd_effShift), otherwise the contact is stiff but undamped.
void mjd_flexContact_mul(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec,
mjtNum scale) {
if (!d->ncon) {
return;
}
int nv = m->nv;
mj_markStack(d);
mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacn = mjSTACKALLOC(d, 3*nv, mjtNum);
int* chain = mjSTACKALLOC(d, nv, int);
for (int i = 0; i < d->ncon; i++) {
const mjContact* con = d->contact + i;
if (con->exclude != 4) {
continue;
}
mjtNum k = mjd_flexContactStiffness(m, d, con);
if (k <= 0) {
continue;
}
int NV = mj_contactJacobian(m, d, con, con->dim, jacdif, NULL, jac1, jac2, NULL, NULL, chain);
if (!NV) {
continue;
}
mju_mulMatMat(jacn, con->frame, jacdif, con->dim > 1 ? 3 : 1, 3, NV);
mjtNum Jv = 0;
for (int a = 0; a < NV; a++) {
Jv += jacn[a] * vec[chain[a]];
}
mjtNum s = scale * k * Jv;
for (int a = 0; a < NV; a++) {
res[chain[a]] += s * jacn[a];
}
}
mj_freeStack(d);
}
mjtBool mjd_flexStiff_any(const mjModel* m, int flg_interp) {
for (int f = 0; f < m->nflex; f++) {
if (flg_interp && flexInterp_processed(m, f)) {
@@ -1645,6 +1698,11 @@ mjtBool mjd_flexStiff_any(const mjModel* m, int flg_interp) {
// does this standard flex contribute implicit stiffness under the given term flags?
// A flex participates if it has elasticity OR passive contacts: the contact stiffness may be the
// only stiffness, so vertex slots must exist either way.
static mjtBool flexMetric_participates(const mjModel* m, int f, int flg_bend, int flg_stretch,
int flg_contact);
static mjtBool flexStiff_active(const mjModel* m, int f, int flg_bend, int flg_stretch) {
if (m->flex_interp[f] || m->flex_rigid[f] || m->flex_dim[f] < 2) {
return 0;
@@ -1655,6 +1713,88 @@ static mjtBool flexStiff_active(const mjModel* m, int f, int flg_bend, int flg_s
return bend || stretch;
}
// Passive contact stiffness: k = omega^2 * m_min, a natural frequency scaled by the smallest
// nonzero participating mass (pinned vertices carry mass 0 and are skipped).
#define mjFLEXCONTACT_OMEGA2 5e7
mjtNum mjd_flexContactStiffness(const mjModel* m, const mjData* d, const mjContact* con) {
mjtNum mmin = 0;
for (int side = 0; side < 2; side++) {
int f = con->flex[side];
if (f < 0) {
continue;
}
int gv[8], ngv = 0;
if (con->vert[side] >= 0) {
gv[ngv++] = m->flex_vertadr[f] + con->vert[side];
} else if (con->elem[side] >= 0) {
int nvrt = m->flex_dim[f] + 1;
const int* e = m->flex_elem + m->flex_elemdataadr[f] + nvrt*con->elem[side];
for (int j = 0; j < nvrt && ngv < 8; j++) {
gv[ngv++] = m->flex_vertadr[f] + e[j];
}
}
for (int j = 0; j < ngv; j++) {
int b = m->flex_vertbodyid[gv[j]];
if (m->body_dofnum[b] != 3) {
continue;
}
int da = m->body_dofadr[b];
mjtNum mv = d->M[m->M_rowadr[da] + m->M_rownnz[da] - 1]; // diagonal: the point mass
if (mv > 0 && (mmin == 0 || mv < mmin)) {
mmin = mv;
}
}
}
return mjFLEXCONTACT_OMEGA2 * mmin; // 0 if every participant is massless: no stiffness, no NaN
}
// The flex vertex slots a passive contact couples: the vertex itself for a vertex side, the
// element's vertices for an element side. Duplicates dropped, and slots outside the metric skipped.
static int contactFlexSlots(const mjModel* m, const mjContact* con, const int* vslot,
int* out, int cap) {
int n = 0;
for (int side = 0; side < 2; side++) {
int f = con->flex[side];
if (f < 0) {
continue;
}
int gv[8], ngv = 0;
if (con->vert[side] >= 0) {
gv[ngv++] = m->flex_vertadr[f] + con->vert[side];
} else if (con->elem[side] >= 0) {
int nvrt = m->flex_dim[f] + 1;
const int* e = m->flex_elem + m->flex_elemdataadr[f] + nvrt*con->elem[side];
for (int j = 0; j < nvrt && ngv < 8; j++) {
gv[ngv++] = m->flex_vertadr[f] + e[j];
}
}
for (int j = 0; j < ngv; j++) {
int s = vslot[gv[j]];
if (s < 0) {
continue;
}
int dup = 0;
for (int q = 0; q < n; q++) {
if (out[q] == s) { dup = 1; break; }
}
if (!dup && n < cap) {
out[n++] = s;
}
}
}
return n;
}
static mjtBool flexMetric_participates(const mjModel* m, int f, int flg_bend, int flg_stretch,
int flg_contact) {
if (flexStiff_active(m, f, flg_bend, flg_stretch)) {
return 1;
}
return flg_contact && m->flex_passive[f] && !m->flex_rigid[f] && !m->flex_interp[f] &&
m->flex_dim[f] >= 2;
}
// assemble the standard-flex implicit stiffness K = (s1 + s2*damping) * (K_bend + K_stretch)
// into dof-level CSR (same terms mjd_flexBend_mul / mjd_flexStretch_mul apply matrix-free; the
@@ -1669,7 +1809,7 @@ static mjtBool flexStiff_active(const mjModel* m, int f, int flg_bend, int flg_s
// so one CSR replaces all three matrix-free operators uniformly.
int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr,
int* colind, mjtNum* val, mjtNum s1, mjtNum s2,
int flg_bend, int flg_stretch, const mjtNum* Krot) {
int flg_bend, int flg_stretch, int flg_contact, const mjtNum* Krot) {
int nv = m->nv;
mj_markStack(d);
@@ -1680,7 +1820,7 @@ int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr
vslot[i] = -1;
}
for (int f = 0; f < m->nflex; f++) {
if (!flexStiff_active(m, f, flg_bend, flg_stretch)) {
if (!flexMetric_participates(m, f, flg_bend, flg_stretch, flg_contact)) {
continue;
}
for (int lv = 0; lv < m->flex_vertnum[f]; lv++) {
@@ -1777,6 +1917,20 @@ int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr
}
}
// passive contacts (counting): each contact makes its vertices mutual neighbours in the CSR.
if (flg_contact) {
for (int i = 0; i < d->ncon; i++) {
const mjContact* con = d->contact + i;
if (con->exclude != 4) {
continue;
}
int cs[8], ncs = contactFlexSlots(m, con, vslot, cs, 8);
for (int a = 0; a < ncs; a++) {
ncand[cs[a]] += ncs;
}
}
}
// gather candidate neighbor lists (vertex slots, with duplicates)
int* cadr = mjSTACKALLOC(d, nvert + 1, int);
cadr[0] = 0;
@@ -1842,6 +1996,22 @@ int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr
}
}
// passive contacts (filling)
if (flg_contact) {
for (int i = 0; i < d->ncon; i++) {
const mjContact* con = d->contact + i;
if (con->exclude != 4) {
continue;
}
int cs[8], ncs = contactFlexSlots(m, con, vslot, cs, 8);
for (int a = 0; a < ncs; a++) {
for (int b = 0; b < ncs; b++) {
cand[cadr[cs[a]] + ncand[cs[a]]++] = cs[b];
}
}
}
}
// per vertex: sort by neighbor dofadr, unique -> neighbor lists
int* nadr = mjSTACKALLOC(d, nvert + 1, int);
int* neigh = mjSTACKALLOC(d, cadr[nvert] > 0 ? cadr[nvert] : 1, int);
@@ -2103,6 +2273,69 @@ int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr
})
}
}
// passive contacts (values): assemble k*J^T*J blocks, where J is the contact-normal Jacobian.
// All participants are metric-carried flex vertices, so the block is assembled in full.
if (flg_contact && d->ncon) {
int* dof2slot = mjSTACKALLOC(d, nv, int);
mjtNum* jacdif = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jac1 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jac2 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacn = mjSTACKALLOC(d, 3*nv, mjtNum);
int* chain = mjSTACKALLOC(d, nv, int);
mjtNum* w = mjSTACKALLOC(d, 3*(nvert > 0 ? nvert : 1), mjtNum);
for (int i = 0; i < nv; i++) {
dof2slot[i] = -1;
}
for (int s = 0; s < nvert; s++) {
for (int k = 0; k < 3; k++) {
dof2slot[vdof[s] + k] = s;
}
}
for (int i = 0; i < d->ncon; i++) {
const mjContact* con = d->contact + i;
if (con->exclude != 4) {
continue;
}
int cs[8], ncs = contactFlexSlots(m, con, vslot, cs, 8);
if (ncs < 1) {
continue;
}
mjtNum k = mjd_flexContactStiffness(m, d, con);
if (k <= 0) {
continue;
}
int NV = mj_contactJacobian(m, d, con, con->dim, jacdif, NULL, jac1, jac2, NULL, NULL, chain);
if (NV == 0) {
continue;
}
// rotate into the contact frame and keep the normal row
mju_mulMatMat(jacn, con->frame, jacdif, con->dim > 1 ? 3 : 1, 3, NV);
for (int a = 0; a < ncs; a++) {
mju_zero(w + 3*cs[a], 3);
}
for (int a = 0; a < NV; a++) {
int s = dof2slot[chain[a]];
if (s >= 0) {
w[3*s + (chain[a] - vdof[s])] = jacn[a];
}
}
for (int a = 0; a < ncs; a++) {
for (int b = 0; b < ncs; b++) {
int pos;
FLEXSTIFF_BLOCK(cs[a], cs[b], pos);
if (pos < 0) {
continue;
}
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
val[rowadr[vdof[cs[a]] + r] + 3*pos + c] += s1 * k * w[3*cs[a] + r] * w[3*cs[b] + c];
}
}
}
}
}
}
#undef FLEXSTIFF_BLOCK
#undef FLEXINTERP_WALK
@@ -3166,6 +3399,7 @@ void mjd_effShift(const mjModel* m, mjData* d) {
mjd_flexInterp_mul(m, d, d->efm_c, d->qvel, h, 0, d->flexelem_krot);
mjd_flexBend_mul(m, d, d->efm_c, d->qvel, -h, 0);
mjd_flexStretch_mul(m, d, d->efm_c, d->qvel, -h, 0);
mjd_flexContact_mul(m, d, d->efm_c, d->qvel, -h);
}
@@ -3197,16 +3431,17 @@ void mjd_effBuild(const mjModel* m, mjData* d, int active, int flg_factor) {
const mjtNum* krot = mjd_flexInterpAssemblable(m) ? d->flexelem_krot : NULL;
d->efm_K_rownnz = EFMALLOC(int, nv);
d->efm_K_rowadr = EFMALLOC(int, nv);
if (mjd_flexStiff_any(m, krot != NULL)) {
if (mjd_flexStiff_any(m, krot != NULL) || flexPassiveContact_any(m)) {
d->nefmK = mjd_flexStiff_assemble(m, d, d->efm_K_rownnz, d->efm_K_rowadr,
NULL, NULL, h*h, h, /*bend*/ 1, /*stretch*/ 1, krot);
NULL, NULL, h*h, h, /*bend*/ 1, /*stretch*/ 1,
/*contact*/ 1, krot);
}
if (d->nefmK) {
d->efm_K_colind = EFMALLOC(int, d->nefmK);
d->efm_K_val = EFMALLOC(mjtNum, d->nefmK);
mjd_flexStiff_assemble(m, d, d->efm_K_rownnz, d->efm_K_rowadr,
d->efm_K_colind, d->efm_K_val, h*h, h,
/*bend*/ 1, /*stretch*/ 1, krot);
/*bend*/ 1, /*stretch*/ 1, /*contact*/ 1, krot);
// per-step factor of the flex block of (M + K): the stiffness is constant during the
// step, so one factorization here turns every preconditioner application into a direct
// solve (the stiff flex block stops being iterated on). Consumers that only multiply
@@ -3219,6 +3454,7 @@ void mjd_effBuild(const mjModel* m, mjData* d, int active, int flg_factor) {
mju_zeroInt(d->efm_K_rownnz, nv);
mju_zeroInt(d->efm_K_rowadr, nv);
}
d->efm_active = 1;
// fill the shift with the current velocity (refreshed again in the velocity stage)
src/engine/engine_derivative.h
+8 -1
View File
@@ -77,9 +77,16 @@ MJAPI void mjd_flexStretch_mul(const mjModel* m, mjData* d, mjtNum* res, const m
// dof-level CSR; phase 1 (colind==NULL) fills rownnz/rowadr and returns total nnz, phase 2
// fills colind/val. Interp flexes are assembled iff Krot (mjd_flexInterp_cacheKrot cache) is
// non-NULL and the centered fast path applies (check mjd_flexInterpAssemblable first).
// Passive contact stiffness (omega^2 * m_min); force and Hessian must use the same value.
// res += scale * K_contact * vec (shift counterpart of the contact stiffness in the metric).
MJAPI void mjd_flexContact_mul(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec,
mjtNum scale);
MJAPI mjtNum mjd_flexContactStiffness(const mjModel* m, const mjData* d, const mjContact* con);
MJAPI int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr,
int* colind, mjtNum* val, mjtNum s1, mjtNum s2,
int flg_bend, int flg_stretch, const mjtNum* Krot);
int flg_bend, int flg_stretch, int flg_contact, const mjtNum* Krot);
// can all interp flexes be assembled to dof-level CSR? (centered fast path everywhere)
MJAPI mjtBool mjd_flexInterpAssemblable(const mjModel* m);
src/engine/engine_forward.c
+17 -1
View File
@@ -1591,6 +1591,16 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
// return 1 if any flex needs implicit stiffness treatment (interp or bending)
// return 1 if any non-rigid flex uses passive contacts (needs the metric independently of elasticity)
static mjtBool flex_has_passive_contact(const mjModel* m) {
for (int f=0; f < m->nflex; f++) {
if (!m->flex_rigid[f] && m->flex_passive[f]) {
return 1;
}
}
return 0;
}
static mjtBool flex_has_implicit_stiffness(const mjModel* m) {
for (int f=0; f < m->nflex; f++) {
if (m->flex_rigid[f]) {
@@ -1634,7 +1644,7 @@ int mj_flexCG(const mjModel* m) {
return m->opt.solver == mjSOL_CG &&
(m->opt.integrator == mjINT_IMPLICIT || m->opt.integrator == mjINT_IMPLICITFAST) &&
m->opt.cone != mjCONE_ELLIPTIC && !mjENABLED(mjENBL_SLEEP) &&
flex_has_implicit_stiffness(m);
(flex_has_implicit_stiffness(m) || flex_has_passive_contact(m));
}
@@ -1776,6 +1786,12 @@ void mj_implicit(const mjModel* m, mjData* d) {
void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) {
TM_START;
// Passive flex contact is too stiff for explicit integration; require the effective metric.
if (flex_has_passive_contact(m) && !mj_flexCG(m)) {
mjERROR("passive flex contact requires the effective metric: use integrator=\"implicit\" or "
"\"implicitfast\" with solver=\"CG\", pyramidal cones and sleep disabled");
}
// position-dependent
if (skipstage < mjSTAGE_POS) {
mj_fwdPosition(m, d);
src/engine/engine_passive.c
+3 -4
View File
@@ -22,6 +22,7 @@
#include "engine/engine_callback.h"
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_util.h"
#include "engine/engine_derivative.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_inline.h"
#include "engine/engine_memory.h"
@@ -36,8 +37,6 @@
//----------------------------- passive forces -----------------------------------------------------
// stiffness for passive contacts
static const mjtNum kContactStiffness = 1e4;
// local edge-based vertex indexing for 2D and 3D elements, 2D and 3D elements
// have 3 and 6 edges, respectively so the missing indexes are set to 0
@@ -962,8 +961,8 @@ int mj_contactPassive(const mjModel* m, mjData* d) {
// rotate Jacobian differences to contact frame
mju_mulMatMat(jac, con->frame, jacdifp, dim > 1 ? 3 : 1, 3, NV);
// compute passive contact force (dim = 1)
mjtNum scl = -kContactStiffness*con->dist;
// compute passive contact force (dim = 1); stiffness shared with the metric Hessian.
mjtNum scl = -mjd_flexContactStiffness(m, d, con)*con->dist;
if (!issparse) {
mju_addToScl(d->qfrc_spring, jac, scl, nv);
} else {
src/engine/engine_setconst.c
+3 -2
View File
@@ -1416,10 +1416,11 @@ static void setEfm0Factor(mjModel* m, mjData* d) {
int* K_rownnz = mjSTACKALLOC(d, nv, int);
int* K_rowadr = mjSTACKALLOC(d, nv, int);
int nK = mjd_flexStiff_assemble(m, d, K_rownnz, K_rowadr, NULL, NULL, h*h, h,
/*flg_bend=*/1, /*flg_stretch=*/0, NULL);
/*flg_bend=*/1, /*flg_stretch=*/0, /*flg_contact=*/0,
NULL);
int* K_colind = mjSTACKALLOC(d, nK > 0 ? nK : 1, int);
mjtNum* K_val = mjSTACKALLOC(d, nK > 0 ? nK : 1, mjtNum);
mjd_flexStiff_assemble(m, d, K_rownnz, K_rowadr, K_colind, K_val, h*h, h, 1, 0, NULL);
mjd_flexStiff_assemble(m, d, K_rownnz, K_rowadr, K_colind, K_val, h*h, h, 1, 0, 0, NULL);
// inverse map: dof address -> compact factor row (monotone: slots follow dof order)
int* dofrow = mjSTACKALLOC(d, nv, int);