From 0041fdcbb0565ee26ea1d7e7af9198eeb5c04b7e Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Mon, 9 Feb 2026 11:53:56 -0800 Subject: [PATCH] Add implicit stiffness for flex_interp to mj_implicitSkip. PiperOrigin-RevId: 867706885 Change-Id: Ic94c65b618a415609bffe3d69a86f9034f2d2400 --- doc/XMLreference.rst | 2 +- doc/changelog.rst | 4 + src/engine/engine_derivative.c | 241 ++++++++++++++++++++ src/engine/engine_derivative.h | 9 + src/engine/engine_forward.c | 158 +++++++++++-- src/engine/engine_passive.c | 110 +++++---- src/engine/engine_passive.h | 4 + test/engine/engine_derivative_test.cc | 313 +++++++++++++++++++++++++- test/engine/engine_forward_test.cc | 139 ++++++++++++ 9 files changed, 910 insertions(+), 70 deletions(-) diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 0787afce..5d08dee3 100644 --- a/doc/XMLreference.rst +++ b/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 `__, diff --git a/doc/changelog.rst b/doc/changelog.rst index d7d456ec..5f1c46cd 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -20,6 +20,10 @@ Significant new features - Added new :ref:`flexvert` equality constraints that enable cloth simulations with coarser meshes. This adds a new option ``vert`` to flexcomp edge :ref:`equality` and the new equality type :ref:`flexvert`. +- Added implicit integration support for deformable objects (flex) in ``implicit`` and ``implicitfast`` + :ref:`integrators`. 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` types. .. container:: custom-clear diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index 7a11c1f7..47552341 100644 --- a/src/engine/engine_derivative.c +++ b/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; inflex; 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= 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 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; diff --git a/src/engine/engine_derivative.h b/src/engine/engine_derivative.h index 8f316e6c..af2aef4c 100644 --- a/src/engine/engine_derivative.h +++ b/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 diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 7d767a13..b00810d5 100644 --- a/src/engine/engine_forward.c +++ b/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; nflex_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; nflex_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 diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 6653b4ca..d52e478c 100644 --- a/src/engine/engine_passive.c +++ b/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; } diff --git a/src/engine/engine_passive.h b/src/engine/engine_passive.h index d1d4d949..ba8b3233 100644 --- a/src/engine/engine_passive.h +++ b/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 ----------------------------------------------------------- diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index c887f4b7..1a050338 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -14,6 +14,9 @@ // Tests for engine/engine_derivative.c. +#include "src/engine/engine_derivative.h" + +#include #include #include #include @@ -23,7 +26,6 @@ #include #include #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"( + + + )"; + + 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 vec(nv); + std::vector 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 H(nv * nv, 0); + std::vector 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 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& 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 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 qDerivAnalytic(nD); + mju_zero(data->qDeriv, nD); + mjd_passive_vel(model, data); + mju_copy(qDerivAnalytic.data(), data->qDeriv, nD); + + // finite-difference derivatives + std::vector 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 dof_indices(nv); + for (int i = 0; i < nv; i++) dof_indices[i] = i; + + vector H1(nv * nv, 0); + mjd_flexInterp_addH(model, data, H1.data(), dof_indices.data(), nv, 1.0); + + vector H2(nv * nv, 0); + mjd_flexInterp_addH(model, data, H2.data(), dof_indices.data(), nv, 0.5); + + vector 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"( + + + )"; + + 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 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 H_approx(nv * nv, 0); + std::vector 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 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 diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc index c4f73c1f..8d070ae2 100644 --- a/test/engine/engine_forward_test.cc +++ b/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 #include @@ -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"( + + + + + + + + + + + + + + + )"; + 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 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"( + + + )"; + 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