Add implicit stiffness for flex_interp to mj_implicitSkip.

PiperOrigin-RevId: 867706885
Change-Id: Ic94c65b618a415609bffe3d69a86f9034f2d2400

doc/XMLreference.rst

@@ -3575,7 +3575,7 @@ saving the XML:
 
 .. _body-flexcomp-dof:
 
-:at:`dof`: :at-val:`[full, radial, trilinear], "full"`
+:at:`dof`: :at-val:`[full, radial, trilinear, quadratic], "full"`
    The parametrization of the flex's degrees of freedom (dofs). See the video on the right illustrating the
    different parametrizations with deformable spheres. The three models in the video are respectively
    `sphere_full <https://github.com/google-deepmind/mujoco/blob/main/model/flex/sphere_full.xml>`__,

doc/changelog.rst

@@ -20,6 +20,10 @@ Significant new features
 - Added new :ref:`flexvert<equality-flexvert>` equality constraints that enable cloth simulations with coarser meshes.
   This adds a new option ``vert`` to flexcomp edge :ref:`equality<flexcomp-edge-equality>` and the new equality type
   :ref:`flexvert<equality-flexvert>`.
+- Added implicit integration support for deformable objects (flex) in ``implicit`` and ``implicitfast``
+  :ref:`integrators<geIntegration>`. This method extracts the flex degrees of freedom and solves them as a dense block,
+  enabling increased stability for stiff flex objects without reducing the timestep. It is compatible with the
+  ``trilinear`` and ``quadratic`` :ref:`dof<body-flexcomp-dof>` types.
 
 .. container:: custom-clear
 

src/engine/engine_derivative.c

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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 -----------------------------------------------------------
 

test/engine/engine_derivative_test.cc

@@ -14,6 +14,9 @@
 
 // Tests for engine/engine_derivative.c.
 
+#include "src/engine/engine_derivative.h"
+
+#include <cstddef>
 #include <random>
 #include <string>
 #include <vector>
@@ -23,7 +26,6 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_core_smooth.h"
-#include "src/engine/engine_derivative.h"
 #include "src/engine/engine_derivative_fd.h"
 #include "src/engine/engine_forward.h"
 #include "src/engine/engine_io.h"
@@ -1078,5 +1080,314 @@ TEST_F(DerivativeTest, quatIntegrate) {
   }
 }
 
+// Utility: Rotate flex grid
+void RotateFlexGrid(mjModel* model, mjData* data, const char* flex_name,
+                    double angle) {
+  int flex_id = mj_name2id(model, mjOBJ_FLEX, flex_name);
+  ASSERT_NE(flex_id, -1);
+  int node_adr = model->flex_nodeadr[flex_id];
+  int* node_bodies = model->flex_nodebodyid + node_adr;
+  int nodenum = model->flex_nodenum[flex_id];
+
+  // Make deterministic quaternion for rotation inside helper
+  mjtNum quat[4] = {1, 0, 0, 0};
+  if (angle != 0) {
+    mjtNum vel[3] = {1, 1, 1};
+    mju_normalize3(vel);
+    mju_quatIntegrate(quat, vel, angle);
+  }
+
+  // reset first to get initial positions
+  mj_resetData(model, data);
+  mj_forward(model, data);  // Compute initial xpos
+
+  // Update qpos
+  for (int i = 0; i < nodenum; i++) {
+    int bodyid = node_bodies[i];
+
+    // Only process nodes with valid bodies (FlexInterpDamping assumes this)
+    if (bodyid >= 0) {
+      mjtNum xpos0[3];
+      mju_copy3(xpos0, data->xpos + 3 * bodyid);  // Initial absolute position
+
+      mjtNum xpos_new[3];
+      mju_rotVecQuat(xpos_new, xpos0, quat);  // Rotate absolute position
+
+      mjtNum delta[3];
+      mju_sub3(delta, xpos_new, xpos0);
+
+      // Find the qpos address for this node/body
+      int jnt = model->body_jntadr[bodyid];
+      if (jnt >= 0) {
+        int qadr = model->jnt_qposadr[jnt];
+        mju_addTo3(data->qpos + qadr, delta);
+      }
+    }
+  }
+}
+
+// compare analytic and fin-diff d_qfrc_passive/d_qvel for flex interp
+// Combined test for verify mjd_flexInterp_mulK (stiffness) and damping
+TEST_F(DerivativeTest, FlexInterpDerivatives) {
+  static const char* const kXml = R"(
+  <mujoco>
+    <option integrator="implicit"/>
+    <worldbody>
+      <flexcomp name="flex" type="grid" count="3 3 3" spacing="0.1 0.2 0.3"
+                radius=".01" dim="3" mass="1" dof="trilinear">
+        <contact selfcollide="none"/>
+        <elasticity young="1e4" poisson="0.3" damping="50"/>
+      </flexcomp>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024];
+  mjModel* model = LoadModelFromString(kXml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  int nD = model->nD;
+  int nv = model->nv;
+  ASSERT_EQ(model->nq, 24);  // 8 corners * 3 dofs
+
+  mjData* data = mj_makeData(model);
+
+  // iterate over rotations
+  for (mjtNum angle : {0.0, 0.5, 1.0, mjPI / 2, mjPI, 2.0 * mjPI}) {
+    RotateFlexGrid(model, data, "flex", angle);
+    mj_forward(model, data);
+
+    // part 1: stiffness verification
+    {
+      std::vector<mjtNum> vec(nv);
+      std::vector<mjtNum> res(nv);
+      mju_zero(vec.data(), nv);
+      // use deterministic random perturbation to verify full stiffness matrix
+      // behavior
+      for (int i = 0; i < nv; i++) {
+        vec[i] = mju_Halton(i, 2) - 0.5;
+      }
+
+      // use addH to compute K * vec
+      // addH adds (h^2*K + h*D) to H
+      // if we set h=1, damping=0, we get K added to H
+      mjtNum save_damping = model->flex_damping[0];
+      model->flex_damping[0] = 0;
+
+      std::vector<mjtNum> H(nv * nv, 0);
+      std::vector<int> dof_indices(nv);
+      for (int i = 0; i < nv; i++) dof_indices[i] = i;
+
+      // assemble K into H
+      mjd_flexInterp_addH(model, data, H.data(), dof_indices.data(), nv, 1.0);
+
+      // restore damping
+      model->flex_damping[0] = save_damping;
+
+      // compute res = K * vec
+      mju_mulMatVec(res.data(), H.data(), vec.data(), nv, nv);
+
+      // finite difference of mj_passive for stiffness
+      double eps = 1e-6;
+      mjData* data_perturbed = mj_copyData(NULL, model, data);
+
+      // apply perturbation
+      mju_addToScl(data_perturbed->qpos, vec.data(), eps, nv);
+
+      // recompute geometry/passive
+      mj_forward(model, data_perturbed);
+
+      // compute FD estimate of K * vec
+      // qfrc_passive = -dV/dq => d(qfrc)/dq = -K
+      // (qfrc_new - qfrc)/eps ~= -K * vec
+      std::vector<mjtNum> fd_res(nv);
+      for (int i = 0; i < nv; ++i) {
+        fd_res[i] =
+            -(data_perturbed->qfrc_passive[i] - data->qfrc_passive[i]) / eps;
+      }
+
+      // compare analytical result (H*vec) with FD result
+      for (int i = 0; i < nv; ++i) {
+        EXPECT_NEAR(res[i], fd_res[i], 5e-3)
+            << "Stiffness Mismatch at DOF " << i;
+      }
+
+      mj_deleteData(data_perturbed);
+
+      // check symmetry: K[i,j] == K[j,i]
+      std::vector<mjtNum>& K_full = H;
+      mjtNum max_asymmetry = 0;
+      for (int i = 0; i < nv; i++) {
+        for (int j = 0; j < i; j++) {
+          mjtNum diff = mju_abs(K_full[i * nv + j] - K_full[j * nv + i]);
+          max_asymmetry = mju_max(max_asymmetry, diff);
+        }
+      }
+      EXPECT_LT(max_asymmetry, 1e-10)
+          << "K matrix is not symmetric at angle " << angle;
+
+      // check positive semi-definiteness: v^T * K * v >= 0
+      for (int trial = 0; trial < 5; trial++) {
+        std::vector<mjtNum> v(nv);
+        for (int i = 0; i < nv; i++) {
+          v[i] = mju_Halton(i + trial * nv, 3) - 0.5;
+        }
+        mjtNum vKv = 0;
+        for (int i = 0; i < nv; i++) {
+          for (int j = 0; j < nv; j++) {
+            vKv += v[i] * K_full[i * nv + j] * v[j];
+          }
+        }
+        EXPECT_GE(vKv, -1e-8) << "K matrix is not PSD at angle " << angle;
+      }
+    }
+
+    // part 2: damping verification
+    {
+      // set velocity non-zero to test damping
+      data->qvel[0] = 1.0;
+
+      mj_forward(model, data);
+
+      // get analytic derivatives (without Flex Damping currently)
+      std::vector<mjtNum> qDerivAnalytic(nD);
+      mju_zero(data->qDeriv, nD);
+      mjd_passive_vel(model, data);
+      mju_copy(qDerivAnalytic.data(), data->qDeriv, nD);
+
+      // finite-difference derivatives
+      std::vector<mjtNum> qDerivFD(nD);
+      mju_zero(data->qDeriv, nD);
+      mjtNum eps = 1e-6;
+
+      mjd_passive_velFD(model, data, eps);
+      mju_copy(qDerivFD.data(), data->qDeriv, nD);
+
+      // check that we have non-zero damping (FD should find it)
+      EXPECT_GT(mju_norm(qDerivFD.data(), nD), 1e-3);
+
+      // compute expected flex damping using mjd_flexInterp_addH
+      // D = 4*H(0.5) - H(1)
+      vector<int> dof_indices(nv);
+      for (int i = 0; i < nv; i++) dof_indices[i] = i;
+
+      vector<mjtNum> H1(nv * nv, 0);
+      mjd_flexInterp_addH(model, data, H1.data(), dof_indices.data(), nv, 1.0);
+
+      vector<mjtNum> H2(nv * nv, 0);
+      mjd_flexInterp_addH(model, data, H2.data(), dof_indices.data(), nv, 0.5);
+
+      vector<mjtNum> D(nv * nv);
+      for (int i = 0; i < nv * nv; i++) {
+        D[i] = 4.0 * H2[i] - H1[i];
+      }
+
+      // subtract D from qDerivAnalytic using sparse indexing
+      // d(force)/d(vel) = -D
+      for (int i = 0; i < nv; i++) {
+        int rownnz = model->D_rownnz[i];
+        int rowadr = model->D_rowadr[i];
+        for (int k = 0; k < rownnz; k++) {
+          int index = rowadr + k;
+          int j = model->D_colind[index];
+          qDerivAnalytic[index] -= D[i * nv + j];
+        }
+      }
+
+      // expect FD and corrected analytic derivatives to match
+      mjtNum tol = 1e-4;
+      EXPECT_THAT(qDerivAnalytic, Pointwise(DoubleNear(tol), qDerivFD))
+          << "Damping Mismatch at angle: " << angle;
+    }
+  }
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+// Test Jacobian under deformation to highlight approximation error
+TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) {
+  static const char* const kXml = R"(
+  <mujoco>
+    <option integrator="implicit"/>
+    <worldbody>
+      <flexcomp name="flex" type="grid" count="3 3 3" spacing="0.1 0.2 0.3"
+                radius=".01" dim="3" mass="1" dof="trilinear">
+        <contact selfcollide="none"/>
+        <elasticity young="1e4" poisson="0.3" damping="0"/>
+      </flexcomp>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024];
+  mjModel* model = LoadModelFromString(kXml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  int nv = model->nv;
+
+  mjData* data = mj_makeData(model);
+
+  // Apply rotation
+  RotateFlexGrid(model, data, "flex", 1.0);  // 1 radian rotation
+
+  // Apply deformation (stretch along X)
+  // qpos is initialized by RotateFlexGrid.
+  // Add a random perturbation to qpos that represents deformation.
+  // We use a deterministic sequence to ensure reproducibility.
+  std::vector<mjtNum> deformation(nv);
+  for (int i = 0; i < nv; i++) {
+    // Large deformation to make sure terms are significant
+    deformation[i] = (mju_Halton(i, 3) - 0.5) * 0.2;
+  }
+  mju_addTo(data->qpos, deformation.data(), nv);
+
+  mj_forward(model, data);
+
+  // 1. Compute Analytic Jacobian (Approximate)
+  // We use mjd_flexInterp_addH to get K_approx
+  std::vector<mjtNum> H_approx(nv * nv, 0);
+  std::vector<int> dof_indices(nv);
+  for (int i = 0; i < nv; i++) dof_indices[i] = i;
+
+  // h=1, damping=0 => adds K to H
+  mjd_flexInterp_addH(model, data, H_approx.data(), dof_indices.data(), nv,
+                      1.0);
+
+  // 2. Compute Finite Difference Jacobian (Ground Truth)
+  // qfrc_passive = -dV/dq
+  // d(qfrc)/dq = -K_true
+  std::vector<mjtNum> K_true(nv * nv, 0);
+  mjtNum eps = 1e-6;
+
+  for (int i = 0; i < nv; i++) {
+    mjData* data_p = mj_copyData(NULL, model, data);
+    data_p->qpos[i] += eps;
+    mj_forward(model, data_p);
+
+    for (int j = 0; j < nv; j++) {
+      // d(force_j)/d(q_i)
+      mjtNum df = data_p->qfrc_passive[j] - data->qfrc_passive[j];
+      // K_true[j, i] = -df/eps
+      K_true[j * nv + i] = -df / eps;
+    }
+    mj_deleteData(data_p);
+  }
+
+  // 3. Compare and check for significant mismatch
+  mjtNum max_error = 0;
+  for (int i = 0; i < nv * nv; i++) {
+    max_error = mju_max(max_error, mju_abs(H_approx[i] - K_true[i]));
+  }
+
+  // We expect significant error because of deformation + rotation.
+  // The missing term (geometric stiffness) is proportional to stress.
+  // We assert that the error is relatively large to confirm the approximation
+  // exists.
+  EXPECT_GT(max_error, 1e-3)
+      << "Jacobian approximation should differ from FD when deformed";
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco

test/engine/engine_forward_test.cc

@@ -15,6 +15,7 @@
 // Tests for engine/engine_forward.c.
 
 #include "src/engine/engine_forward.h"
+#include "src/engine/engine_derivative.h"
 
 #include <cmath>
 #include <cstdlib>
@@ -1617,5 +1618,143 @@ TEST_F(ForwardTest, ActuatorDelayLinearInterp) {
   mj_deleteModel(model);
 }
 
+TEST_F(ForwardTest, FlexTrilinearInstability) {
+  // model parameters matches user's trilinear.xml
+  constexpr char xml[] = R"(
+  <mujoco model="stability_test">
+      <option gravity="0 0 -9.81" iterations="100" solver="CG" tolerance="1e-10"
+              timestep="0.002" integrator="implicitfast">
+          <flag warmstart="disable" island="disable"/>
+      </option>
+      <worldbody>
+          <geom name="floor" size="0 0 .05" type="plane" condim="3"/>
+          <flexcomp name="bed" type="grid" count="17 17 3" spacing="0.05 0.05 0.05"
+                    pos="0 0 0.05" radius="0.0005" dim="3" mass="10" dof="trilinear">
+              <contact condim="3" solref="0.005 1" solimp=".99 .99 .001" selfcollide="none"/>
+              <elasticity young="865067.00" poisson="0.1" damping="1"/>
+          </flexcomp>
+          <body name="box" pos="0.05 0.05 0.5">
+              <freejoint/>
+              <geom name="box_geom" type="box" size="0.04 0.04 0.04" mass="0.5"
+                    solref="0.001 1" solimp="0.99 0.99 0.01"/>
+          </body>
+      </worldbody>
+  </mujoco>
+  )";
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+
+  mjData* data = mj_makeData(model);
+
+  // flex stiffness sign checks
+  // verify correct sign of flex stiffness derivatives before simulation
+  int nv = model->nv;
+  mjtNum h = model->opt.timestep;
+
+  // create a test vector
+  std::vector<mjtNum> v(nv), Mv(nv), flex_Kv(nv);
+  for (int i = 0; i < nv; i++) v[i] = mju_Halton(i, 2) - 0.5;
+  mjtNum vnorm = mju_norm(v.data(), nv);
+  for (int i = 0; i < nv; i++) v[i] /= vnorm;
+
+  mj_forward(model, data);
+
+  // compute M*v and stiffness contributions
+  mj_mulM(model, data, Mv.data(), v.data());
+
+  // note: we use mjd_flexInterp_mulK here (unscaled by h^2) to check raw
+  // stiffness logic similar to what we expect in the solver now
+  mjtNum* v_copy = (mjtNum*)mju_malloc(nv * sizeof(mjtNum));
+  mju_copy(v_copy, v.data(), nv);
+  mju_zero(flex_Kv.data(), nv);
+
+  // using mulKD for legacy check consistency, but we know it applies h^2+h*d
+  // scaling; actually, let's stick to the high-level property checks from
+  // FlexStiffnessSign which used mulKD
+  mjd_flexInterp_mulKD(model, data, flex_Kv.data(), v.data(), h);
+
+  // compute v^T*M*v and v^T*scale*K*v
+  mjtNum vMv = mju_dot(v.data(), Mv.data(), nv);
+  // mulKD returns -scale*K*v, so -flex_Kv = +scale*K*v
+  mjtNum vKv = -mju_dot(v.data(), flex_Kv.data(), nv);
+
+  // assertions from FlexStiffnessSign
+  EXPECT_GT(vKv, 0) << "Stiffness contribution should be positive";
+  EXPECT_GT(vMv + vKv, vMv) << "Full Hessian should exceed M alone";
+
+  mju_free(v_copy);
+
+  // stability simulation
+  // run for steps to catch instability
+  for (int i = 0; i < 2000; ++i) {
+    mj_step(model, data);
+
+    for (int j = 0; j < model->nq; ++j) {
+      if (mju_abs(data->qpos[j]) > 1000.0) {
+        ADD_FAILURE() << "Instability detected at step " << i << " dof " << j
+                      << " val " << data->qpos[j];
+        return;  // Exit early
+      }
+    }
+  }
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+// Verify that flex damping does not affect rigid body motion
+TEST_F(ForwardTest, FlexDampingRigidMotion) {
+  constexpr char xml[] = R"(
+  <mujoco>
+      <option gravity="0 0 0" timestep="0.01" integrator="implicitfast"/>
+      <worldbody>
+          <flexcomp name="flex" type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
+                    pos="0 0 0" euler="45 45 45" radius="0.01" dim="3" mass="1" dof="trilinear">
+              <contact selfcollide="none"/>
+              <elasticity young="1e5" poisson="0.3" damping="10"/>
+          </flexcomp>
+      </worldbody>
+  </mujoco>
+  )";
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  mjData* data = mj_makeData(model);
+
+  // Set initial rigid rotation velocity about Z axis
+  // Center of mass is roughly at 0 0 0 because pos="0 0 0" and symmetric grid.
+  // v = w x r. Let w = (1, 1, 1).
+  mjtNum w[3] = {10.0, 10.0, 10.0};
+  for (int i = 0; i < model->nv / 3; ++i) {
+    int qpos_adr = model->jnt_qposadr[i];
+    int qvel_adr = model->jnt_dofadr[i];
+    mjtNum* pos = data->qpos + qpos_adr;
+    mjtNum* vel = data->qvel + qvel_adr;
+
+    mjtNum r[3] = {pos[0], pos[1], pos[2]};
+    mju_cross(vel, w, r);
+  }
+
+  mj_forward(model, data);
+  mjtNum initial_energy = data->energy[0] + data->energy[1];
+
+  // Run a few steps
+  for (int i = 0; i < 10; ++i) {
+    mj_step(model, data);
+  }
+
+  mj_forward(model, data);
+  mjtNum final_energy = data->energy[0] + data->energy[1];
+
+  // Expect energy conservation.
+  // With the bug, damping force acts on rigid rotation, dissipating energy.
+  EXPECT_NEAR(final_energy, initial_energy, 1e-6 * initial_energy)
+      << "Energy decayed significantly (" << initial_energy << " -> "
+      << final_energy << ")";
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco