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

Add implicit stiffness for flex_interp to mj_implicitSkip.

Browse Source 
PiperOrigin-RevId: 867706885
Change-Id: Ic94c65b618a415609bffe3d69a86f9034f2d2400
This commit is contained in:
Alessio Quaglino
2026-02-09 11:53:56 -08:00
committed by Copybara-Service
parent c1b3b3063e
commit 0041fdcbb0
9 changed files with 910 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_derivative.c
+241
View File
@@ -826,6 +826,247 @@ static mjtNum mjd_muscleGain_vel(mjtNum len, mjtNum vel, const mjtNum lengthrang
}
//--------------------- utility functions for (d force / d pos) * vec Jacobians --------------------
// add J'*B*J*vec to res, sparse version
static void addJTBJ_mulSparse(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec,
const int* J_rownnz, const int* J_rowadr, const int* J_colind,
const mjtNum* J, const mjtNum* B, int n) {
// allocate temp vectors
mj_markStack(d);
mjtNum* Jv = mjSTACKALLOC(d, n, mjtNum);
mjtNum* BJv = mjSTACKALLOC(d, n, mjtNum);
// Jv = J*vec (Sparse Matrix-Vector Multiplication)
mju_zero(Jv, n);
for (int i=0; i < n; i++) {
int nnz = J_rownnz[i];
int adr = J_rowadr[i];
for (int k=0; k < nnz; k++) {
Jv[i] += J[adr + k] * vec[J_colind[adr + k]];
}
}
// BJv = B*Jv (Dense Matrix-Vector Multiplication)
mju_mulMatVec(BJv, B, Jv, n, n);
// res += J'*BJv (Sparse Transpose Matrix-Vector Multiplication)
for (int i=0; i < n; i++) {
int nnz = J_rownnz[i];
int adr = J_rowadr[i];
mjtNum val = BJv[i];
for (int k=0; k < nnz; k++) {
res[J_colind[adr + k]] += J[adr + k] * val;
}
}
mj_freeStack(d);
}
// operation type for flex interpolation derivative kernel
typedef enum {
mjFLEXOP_VEC, // res += J'*K*J*vec
mjFLEXOP_ADDH // H -= J'*K*J to H (dense)
} mjtFlexOp;
// shared kernel for flex interpolation derivatives, scale = s1 + s2*damping
// op: operation type (VEC, or ADDH)
// res: output vector (VEC) or dense H matrix (ADDH)
// vec: input vector for VEC operation, NULL otherwise
// dof_indices, ndof: DOF mapping for ADDH, ignored otherwise
static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
mjtNum* res, const mjtNum* vec, mjtNum s1, mjtNum s2,
const int* dof_indices, int ndof) {
int nv = m->nv;
// build global2local map for ADDH
int* global2local = NULL;
if (op == mjFLEXOP_ADDH) {
mj_markStack(d);
global2local = mjSTACKALLOC(d, nv, int);
mju_fillInt(global2local, -1, nv);
for (int i=0; i<ndof; i++) {
global2local[dof_indices[i]] = i;
}
}
// loop over flexes
for (int f=0; f < m->nflex; f++) {
// only process flex_interp
if (!m->flex_interp[f]) {
continue;
}
// get stiffness and damping
mjtNum* k = m->flex_stiffness + 21*m->flex_elemadr[f];
// skip if rigid or no stiffness
if (m->flex_rigid[f] || k[0] == 0) {
continue;
}
// compute scale
mjtNum damping = m->flex_damping[f];
mjtNum scale = s1 + s2 * damping;
// skip if scale is zero
if (scale == 0) {
continue;
}
int nodenum = m->flex_nodenum[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
// standard stack allocation
mj_markStack(d);
mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* K_rot = mjSTACKALLOC(d, 9*nodenum*nodenum, mjtNum);
// sparse Jacobian allocations
int dim = 3 * nodenum;
int* rownnz = mjSTACKALLOC(d, dim, int);
int* rowadr = mjSTACKALLOC(d, dim, int);
mjtNum* J_val = mjSTACKALLOC(d, dim*nv, mjtNum);
int* J_colind = mjSTACKALLOC(d, dim*nv, int);
// temp allocations for chain
int* chain_colind = mjSTACKALLOC(d, nv, int);
mjtNum* blk_jac = mjSTACKALLOC(d, 3*nv, mjtNum);
// compute positions, rotation and Jacobian
mjtNum quat[4] = {1, 0, 0, 0};
mj_flexInterpState(m, d, f, xpos, NULL, quat);
// compute generalized stiffness in global frame: K_rot = R * K * R^T
mjtNum R[9];
mju_quat2Mat(R, quat); // R = R_global2local
mjtNum RT[9];
mju_transpose(RT, R, 3, 3); // RT = R_local2global
// blockwise rotation: K_rot(i,j) = scale * RT * K_local(i,j) * R
// note: k stores -K, so K_rot = scale * (-K_phys)
for (int i=0; i < nodenum; i++) {
for (int j=0; j < nodenum; j++) {
mjtNum blk[9], tmp[9];
// get K_local(i,j)
int adr = (3*i)*(3*nodenum) + 3*j;
for (int r=0; r < 3; r++) {
for (int c=0; c < 3; c++) {
blk[3*r+c] = k[adr + r*(3*nodenum) + c];
}
}
// tmp = K * R
mju_mulMatMat3(tmp, blk, R);
// blk = RT * tmp = RT * K * R
mju_mulMatMat3(blk, RT, tmp);
// store scaled into K_rot
for (int r=0; r < 3; r++) {
for (int c=0; c < 3; c++) {
K_rot[adr + r*(3*nodenum) + c] = scale * blk[3*r+c];
}
}
}
}
// construct sparse Jacobian J_val
int current_adr = 0;
for (int i=0; i < nodenum; i++) {
// get chain for this node
int chain_nnz = mj_bodyChain(m, bodyid[i], chain_colind);
// compute sparse Jacobian for this node (3 rows)
mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind);
// copy to sparse structure
for (int r=0; r<3; r++) {
int row_idx = 3*i + r;
rownnz[row_idx] = chain_nnz;
rowadr[row_idx] = current_adr;
for (int idx=0; idx<chain_nnz; idx++) {
J_colind[current_adr] = chain_colind[idx];
J_val[current_adr] = blk_jac[r*chain_nnz + idx];
current_adr++;
}
}
}
// perform operation
if (op == mjFLEXOP_VEC) {
// res += J^T * K_rot * J * vec
addJTBJ_mulSparse(m, d, res, vec, rownnz, rowadr, J_colind, J_val, K_rot, dim);
} else if (op == mjFLEXOP_ADDH) {
// H += -J^T * K_rot * J
// H is dense ndof x ndof
// reuse stack for J_reduced (but now we extract from sparse J)
mjtNum* J_reduced = mjSTACKALLOC(d, dim*ndof, mjtNum);
mju_zero(J_reduced, dim*ndof);
// extract columns of J into J_reduced
for (int i=0; i<dim; i++) {
int nnz = rownnz[i];
int adr = rowadr[i];
for (int idx=0; idx<nnz; idx++) {
int global_col = J_colind[adr + idx];
int local_idx = global2local[global_col];
if (local_idx >= 0) {
J_reduced[i*ndof + local_idx] = J_val[adr + idx];
}
}
}
// H -= J_reduced^T * K_rot * J_reduced
// K_rot * J_reduced (dim x ndof)
mjtNum* KJ = mjSTACKALLOC(d, dim*ndof, mjtNum);
mju_mulMatMat(KJ, K_rot, J_reduced, dim, dim, ndof);
// H[i, j] -= sum_k J_reduced[k, i] * KJ[k, j]
for (int i=0; i<ndof; i++) {
for (int j=0; j<ndof; j++) {
mjtNum val = 0;
for (int dim_idx=0; dim_idx<dim; dim_idx++) {
val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j];
}
// res is H
res[i*ndof + j] -= val;
}
}
}
mj_freeStack(d);
}
if (op == mjFLEXOP_ADDH) {
mj_freeStack(d); // free global2local
}
}
// compute res += (h^2 + h*damping) * J'*K*J * vec, for all interpolated flexes
void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, mjtNum h) {
// s1=h*h, s2=h => scale = h*h + h*damping
mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, h * h, h, NULL, 0);
}
// add (h^2 + h*damping) * J'*K*J to dense matrix H, for all interpolated flexes
// H: dense ndof x ndof matrix
// dof_indices: maps local indices to global DOFs
void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, mjtNum h) {
mjd_flexInterp_kernel(m, d, mjFLEXOP_ADDH, H, NULL, h * h, h, dof_indices, ndof);
}
// add (d qfrc_actuator / d qvel) to qDeriv
void mjd_actuator_vel(const mjModel* m, mjData* d) {
int nu = m->nu;
src/engine/engine_derivative.h
+9
View File
@@ -43,6 +43,15 @@ MJAPI void mjd_passive_vel(const mjModel* m, mjData* d);
// subtract (d qfrc_bias / d qvel) from qDeriv (dense version)
MJAPI void mjd_rne_vel_dense(const mjModel* m, mjData* d);
// derivative of flex_interp generalized force w.r.t position: res = (d qfrc_flexinterp / d qpos) * vec
// res and vec are vectors of size m->nv
MJAPI void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, mjtNum h);
// assemble flex stiffness matrix H_flex: H += h*h*K + h*D
// H is a dense matrix of size ndof x ndof, dof_indices maps local rows/cols to global DOFs
MJAPI void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, mjtNum h);
#ifdef __cplusplus
}
#endif
src/engine/engine_forward.c
+135 -23
View File
@@ -1171,9 +1171,24 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
}
// IMPLICIT
if (m->opt.integrator == mjINT_IMPLICIT) {
if (!skipfactor) {
// check for flex_interp
int has_flex_interp = 0;
for (int f = 0; f < m->nflex; f++) {
if (m->flex_interp[f]) {
has_flex_interp = 1;
break;
}
}
// flex: data structures for reduced dense factorization
mjtNum* H_flex = NULL;
int* flex_dof_indices = NULL;
int nflexdofs = 0;
// factorization
if (!skipfactor) {
// implicit
if (m->opt.integrator == mjINT_IMPLICIT) {
// compute analytical derivative qDeriv
mjd_smooth_vel(m, d, /* flg_bias = */ 1);
@@ -1182,20 +1197,10 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
// set qLU = M - dt*qDeriv
mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, nD);
// factorize qLU
int* scratch = mjSTACKALLOC(d, nv, int);
mju_factorLUSparse(d->qLU, nv, scratch, m->D_rownnz, m->D_rowadr, m->D_colind, dof_awake_ind);
}
// solve for qacc: (M - dt*qDeriv) * qacc = qfrc
mju_solveLUSparse(qacc, d->qLU, qfrc, nv, m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind,
dof_awake_ind);
}
// IMPLICITFAST
else if (m->opt.integrator == mjINT_IMPLICITFAST) {
if (!skipfactor) {
// implicitfast
else if (m->opt.integrator == mjINT_IMPLICITFAST) {
// compute analytical derivative qDeriv; skip rne derivative
mjd_smooth_vel(m, d, /* flg_bias = */ 0);
@@ -1204,22 +1209,129 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
// set qH = M - dt*qDeriv
mju_addScl(d->qH, d->M, d->qH, -m->opt.timestep, nC);
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
} else {
mjERROR("integrator must be implicit or implicitfast");
}
// solve for qacc: (M - dt*qDeriv) * qacc = qfrc
// flex: reduced dense factorization
if (has_flex_interp && !sleep_filter) {
// identify flex DOFs
for (int f=0; f < m->nflex; f++) {
if (m->flex_interp[f]) {
int nodenum = m->flex_nodenum[f];
int nodeadr = m->flex_nodeadr[f];
for (int n=0; n<nodenum; n++) {
int b = m->flex_nodebodyid[nodeadr + n];
nflexdofs += m->body_dofnum[b];
}
}
}
// allocations
if (nflexdofs > 0) {
flex_dof_indices = mjSTACKALLOC(d, nflexdofs, int);
int* global2local = mjSTACKALLOC(d, nv, int);
mju_fillInt(global2local, -1, nv);
int cnt = 0;
for (int f=0; f < m->nflex; f++) {
if (m->flex_interp[f]) {
int nodenum = m->flex_nodenum[f];
int nodeadr = m->flex_nodeadr[f];
for (int n=0; n<nodenum; n++) {
int b = m->flex_nodebodyid[nodeadr + n];
int dofnum = m->body_dofnum[b];
int dofadr = m->body_dofadr[b];
for (int j=0; j < dofnum; j++) {
flex_dof_indices[cnt] = dofadr + j;
global2local[dofadr + j] = cnt;
cnt++;
}
}
}
}
// build H_flex (dense) from qLU (implicit) or qH (implicitfast)
H_flex = mjSTACKALLOC(d, nflexdofs*nflexdofs, mjtNum);
mju_zero(H_flex, nflexdofs*nflexdofs);
const int* rownnz = (m->opt.integrator == mjINT_IMPLICIT) ? m->D_rownnz : m->M_rownnz;
const int* rowadr = (m->opt.integrator == mjINT_IMPLICIT) ? m->D_rowadr : m->M_rowadr;
const int* colind = (m->opt.integrator == mjINT_IMPLICIT) ? m->D_colind : m->M_colind;
const mjtNum* source = (m->opt.integrator == mjINT_IMPLICIT) ? d->qLU : d->qH;
for (int i=0; i < nflexdofs; i++) {
int row = flex_dof_indices[i];
int start = rowadr[row];
int end = start + rownnz[row];
for (int k=start; k < end; k++) {
int col = colind[k];
int local_j = global2local[col];
if (local_j >= 0) {
H_flex[i*nflexdofs + local_j] = source[k];
}
}
}
// add stiffness to H_flex
mjtNum h = m->opt.timestep;
mjd_flexInterp_addH(m, d, H_flex, flex_dof_indices, nflexdofs, h);
// factor H_flex
mju_cholFactor(H_flex, nflexdofs, mjMINVAL);
}
}
// standard factorization (implicit / implicitfast)
if (m->opt.integrator == mjINT_IMPLICIT) {
int* scratch = mjSTACKALLOC(d, nv, int);
mju_factorLUSparse(d->qLU, nv, scratch, m->D_rownnz, m->D_rowadr, m->D_colind, dof_awake_ind);
} else {
mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
}
}
// solve
// standard sparse solve
if (m->opt.integrator == mjINT_IMPLICIT) {
mju_solveLUSparse(qacc, d->qLU, qfrc, nv, m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind,
dof_awake_ind);
} else {
// implicitfast
if (sleep_filter) {
mju_copyInd(qacc, qfrc, dof_awake_ind, nv);
} else {
mju_copy(qacc, qfrc, nv);
}
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
}
} else {
mjERROR("integrator must be implicit or implicitfast");
// flex: reduced dense solve
if (H_flex) {
// compute qfrc_flex
mjtNum* qfrc_flex = mjSTACKALLOC(d, nflexdofs, mjtNum);
mjtNum* res = mjSTACKALLOC(d, nv, mjtNum);
mjtNum h = m->opt.timestep;
mjtNum damp = (m->nflex > 0 && m->flex_damping) ? m->flex_damping[0] : 0;
mjtNum scl = h * h + h * damp;
mjtNum factor = (scl > mjMINVAL) ? (h/scl) : 0;
// velocity correction: -h * K * v
mju_zero(res, nv);
mjd_flexInterp_mulKD(m, d, res, d->qvel, h); // returns -scl * K * v
for (int i=0; i < nflexdofs; i++) {
int global_dof = flex_dof_indices[i];
qfrc_flex[i] = qfrc[global_dof] + res[global_dof] * factor;
}
// solve H_flex * qacc_flex = qfrc_flex
// reuse qfrc_flex as result buffer (qacc_flex)
mju_cholSolve(qfrc_flex, H_flex, qfrc_flex, nflexdofs);
// overwrite flex DOFs with reduced dense solution
mju_scatter(qacc, qfrc_flex, flex_dof_indices, nflexdofs);
}
// advance state and time
src/engine/engine_passive.c
+65 -45
View File
@@ -58,6 +58,62 @@ static void inline GradSquaredLengths(mjtNum gradient[6][2][3],
}
}
// compute interpolated flex state: xpos, vel, quat
// f: flex index
// xpos: (output) 3*nodenum
// vel: (output) 3*nodenum, can be NULL
// quat: (output) 4, rotation from global to local
void mj_flexInterpState(const mjModel* m, mjData* d, int f,
mjtNum* xpos, mjtNum* vel, mjtNum* quat) {
int nodenum = m->flex_nodenum[f];
int nstart = m->flex_nodeadr[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
mjtNum com[3] = {0};
// compute positions
if (m->flex_centered[f]) {
for (int i=0; i < nodenum; i++) {
mji_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
if (vel) {
mji_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]);
}
}
} else {
mjtNum screw[6];
for (int i=0; i < nodenum; i++) {
mji_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart));
mji_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
if (vel) {
mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0);
mji_copy3(vel + 3*i, screw + 3);
}
}
}
// compute center of mass
for (int i = 0; i < nodenum; i++) {
mji_addToScl3(com, xpos+3*i, 1.0/nodenum);
}
// compute the Jacobian at the center of mass
mjtNum mat[9] = {0};
mjtNum p[3] = {.5, .5, .5};
mju_defGradient(mat, p, xpos, m->flex_interp[f]);
// find rotation
mju_mat2Rot(quat, mat);
mju_negQuat(quat, quat);
// rotate vertices to quat and add reference center of mass
for (int i = 0; i < nodenum; i++) {
mju_rotVecQuat(xpos+3*i, xpos+3*i, quat);
mji_addTo3(xpos+3*i, p);
if (vel) {
mju_rotVecQuat(vel+3*i, vel+3*i, quat);
}
}
}
// spring and damper forces
static void mj_springdamper(const mjModel* m, mjData* d) {
int nv = m->nv, ntendon = m->ntendon;
@@ -217,55 +273,17 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
}
if (m->flex_interp[f]) {
mjtNum xpos[3*mjMAXFLEXNODES], displ[3*mjMAXFLEXNODES], vel[3*mjMAXFLEXNODES];
mjtNum frc[3*mjMAXFLEXNODES], dmp[3*mjMAXFLEXNODES];
mjtNum com[3] = {0};
mj_markStack(d);
mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* displ = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* vel = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* frc = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* dmp = mjSTACKALLOC(d, 3*nodenum, mjtNum);
mjtNum* xpos0 = m->flex_node0 + 3*m->flex_nodeadr[f];
int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
int nstart = m->flex_nodeadr[f];
// compute positions
if (m->flex_centered[f]) {
for (int i=0; i < nodenum; i++) {
mji_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
mji_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]);
}
} else {
mjtNum screw[6];
for (int i=0; i < nodenum; i++) {
mji_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart));
mji_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0);
mji_copy3(vel + 3*i, screw + 3);
}
}
// compute center of mass
for (int i = 0; i < nodenum; i++) {
mji_addToScl3(com, xpos+3*i, 1.0/nodenum);
}
// re-center positions using center of mass
for (int i = 0; i < nodenum; i++) {
mji_addToScl3(xpos+3*i, com, -1);
}
// compute the Jacobian at the center of mass
mjtNum mat[9] = {0};
mjtNum p[3] = {.5, .5, .5};
mju_defGradient(mat, p, xpos, m->flex_interp[f]);
// find rotation
mjtNum quat[4] = {1, 0, 0, 0};
mju_mat2Rot(quat, mat);
mju_negQuat(quat, quat);
// rotate vertices to quat and add reference center of mass
for (int i = 0; i < nodenum; i++) {
mju_rotVecQuat(xpos+3*i, xpos+3*i, quat);
mji_addTo3(xpos+3*i, p);
mju_rotVecQuat(vel+3*i, vel+3*i, quat);
}
mj_flexInterpState(m, d, f, xpos, vel, quat);
// compute displacement
for (int i = 0; i < nodenum; i++) {
@@ -294,6 +312,8 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
}
}
mj_freeStack(d);
// do not continue with the rest of the flex passive forces
continue;
}
src/engine/engine_passive.h
+4
View File
@@ -28,6 +28,10 @@ extern "C" {
// all passive forces
MJAPI void mj_passive(const mjModel* m, mjData* d);
// compute interpolated flex state: xpos, vel, quat
MJAPI void mj_flexInterpState(const mjModel* m, mjData* d, int f,
mjtNum* xpos, mjtNum* vel, mjtNum* quat);
//------------------------- fluid models -----------------------------------------------------------