From f9f1db1e0a176da2c398458f3b17cd2c47527b38 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Mon, 11 May 2026 10:09:37 -0700
Subject: [PATCH] Replace the banded Cholesky solver for implicit flex
 interpolation with a preconditioned Conjugate Gradient (CG) solver that
 operates on the full system matrix.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The previous approach extracted flex DOFs into a reduced banded system,
factored it separately, and overwrote the global solve. This required
precomputed bandwidth (makeFlexBandwidth), parent-joint detection,
coupling corrections, and a FlexInterpContext struct — and only worked
for standalone flex trees without parent joints.

The new CG solver uses the already-factored global system (M - h*qDeriv)
as a preconditioner and adds the flex stiffness contribution via
matrix-free products (mjd_flexInterp_mulKD/mulK). This handles any
kinematic configuration — including flexes attached to articulated
chains or with parent joints — without sparsity pattern restrictions.

Before (`bunny_multicell`):
```
 Simulation time      : 50.80 s
 Steps per second     : 197
 Realtime factor      : 0.20 x
 Time per step        : 5080.3 µs

 CG iters / step      : 3.16
 Contacts / step      : 31.04
 Constraints / step   : 124.15
 Degrees of freedom   : 178
 Dynamic memory usage : 0.4% of 100M
```

After:
```
 Simulation time      : 9.52 s
 Steps per second     : 1051
 Realtime factor      : 1.05 x
 Time per step        : 951.7 µs

 CG iters / step      : 3.21
 Contacts / step      : 30.90
 Constraints / step   : 123.61
 Degrees of freedom   : 178
 Dynamic memory usage : 0.3% of 100M
```

PiperOrigin-RevId: 913758038
Change-Id: If5aa617b2d535c86aec9bd71c9e0003a2b38bdd7
---
 doc/includes/references.h             |   1 -
 include/mujoco/mjmodel.h              |   1 -
 include/mujoco/mjxmacro.h             |   1 -
 model/flex/bunny_multicell.xml        |   4 +-
 python/mujoco/introspect/structs.py   |   8 -
 src/engine/engine_derivative.c        |  12 +-
 src/engine/engine_derivative.h        |   8 +-
 src/engine/engine_forward.c           | 347 ++++++++------------------
 src/engine/engine_setconst.c          | 131 +---------
 test/engine/engine_derivative_test.cc |  56 +++--
 unity/Runtime/Bindings/MjBindings.cs  |   1 -
 wasm/codegen/generated/bindings.cc    |   4 -
 12 files changed, 152 insertions(+), 422 deletions(-)

diff --git a/doc/includes/references.h b/doc/includes/references.h
index d106b4b6..c226304b 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -1316,7 +1316,6 @@ struct mjModel_ {
   int*      flex_matid;           // material id for rendering                (nflex x 1)
   int*      flex_group;           // group for visibility                     (nflex x 1)
   int*      flex_interp;          // interpolation (0: vertex, 1: nodes)      (nflex x 1)
-  int*      flex_bandwidth;       // precomputed solver bandwidth             (nflex x 1)
   int*      flex_cellnum;         // finite cell num per dimension            (nflex x 3)
   int*      flex_nodeadr;         // first node address                       (nflex x 1)
   int*      flex_nodenum;         // number of nodes                          (nflex x 1)
diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h
index ec2c8250..565358d0 100644
--- a/include/mujoco/mjmodel.h
+++ b/include/mujoco/mjmodel.h
@@ -979,7 +979,6 @@ struct mjModel_ {
   int*      flex_matid;           // material id for rendering                (nflex x 1)
   int*      flex_group;           // group for visibility                     (nflex x 1)
   int*      flex_interp;          // interpolation (0: vertex, 1: nodes)      (nflex x 1)
-  int*      flex_bandwidth;       // precomputed solver bandwidth             (nflex x 1)
   int*      flex_cellnum;         // finite cell num per dimension            (nflex x 3)
   int*      flex_nodeadr;         // first node address                       (nflex x 1)
   int*      flex_nodenum;         // number of nodes                          (nflex x 1)
diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
index 080cbc07..c11ff844 100644
--- a/include/mujoco/mjxmacro.h
+++ b/include/mujoco/mjxmacro.h
@@ -455,7 +455,6 @@
     X   ( int,     flex_matid,            nflex,         1                    ) \
     X   ( int,     flex_group,            nflex,         1                    ) \
     X   ( int,     flex_interp,           nflex,         1                    ) \
-    X   ( int,     flex_bandwidth,        nflex,         1                    ) \
     X   ( int,     flex_cellnum,          nflex,         3                    ) \
     X   ( int,     flex_nodeadr,          nflex,         1                    ) \
     X   ( int,     flex_nodenum,          nflex,         1                    ) \
diff --git a/model/flex/bunny_multicell.xml b/model/flex/bunny_multicell.xml
index 730cbc5d..5f4d9f1d 100644
--- a/model/flex/bunny_multicell.xml
+++ b/model/flex/bunny_multicell.xml
@@ -16,9 +16,9 @@
 <mujoco model="Trilinear">
   <include file="scene.xml"/>
 
-  <option solver="CG" tolerance="1e-6" timestep=".001" integrator="Euler"/>
+  <option solver="CG" tolerance="1e-6" integrator="implicitfast"/>
 
-  <size memory="100M"/>
+  <size memory="10M"/>
 
   <visual>
     <map stiffness="100"/>
diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py
index 84ff1341..5ed73b61 100644
--- a/python/mujoco/introspect/structs.py
+++ b/python/mujoco/introspect/structs.py
@@ -2673,14 +2673,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='interpolation (0: vertex, 1: nodes)',
                  array_extent=('nflex',),
              ),
-             StructFieldDecl(
-                 name='flex_bandwidth',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='precomputed solver bandwidth',
-                 array_extent=('nflex',),
-             ),
              StructFieldDecl(
                  name='flex_cellnum',
                  type=PointerType(
diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c
index ef6bc8b7..85f91e33 100644
--- a/src/engine/engine_derivative.c
+++ b/src/engine/engine_derivative.c
@@ -1134,18 +1134,18 @@ void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const mjtNum
 }
 
 
-// add (h^2 + h*damping) * J'*K*J to banded matrix H, for all interpolated flexes
-//  H: banded ndof x nband matrix (lower triangle, band storage)
-//  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, int nband, mjtNum h) {
-  mjd_flexInterp_kernel(m, d, mjFLEXOP_ADDH, H, NULL, h * h, h, dof_indices, ndof, nband);
+// compute res += h * J'*K*J * vec, for all interpolated flexes (stiffness only, no damping)
+void mjd_flexInterp_mulK(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, mjtNum h) {
+  // s1=h, s2=0 => scale = h (no damping contribution)
+  mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, h, 0, NULL, 0, 0);
 }
 
 
 
 
 
+
+
 // 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 65a136f5..1ddb3af9 100644
--- a/src/engine/engine_derivative.h
+++ b/src/engine/engine_derivative.h
@@ -47,9 +47,11 @@ MJAPI void mjd_rne_vel_dense(const mjModel* m, mjData* d);
 //  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, int nband, mjtNum h);
+// derivative of flex_interp generalized force w.r.t position (stiffness only, no damping)
+MJAPI void mjd_flexInterp_mulK(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, mjtNum h);
+
+
+
 
 
 #ifdef __cplusplus
diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
index ef15c37a..5e285f39 100644
--- a/src/engine/engine_forward.c
+++ b/src/engine/engine_forward.c
@@ -1371,236 +1371,121 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
 }
 
 
-// return 1 if flex f needs implicit interp treatment
-static int flexInterp_active(const mjModel* m, int f) {
-  return m->flex_interp[f] && !m->flex_rigid[f] &&
-         m->flex_edgeequality[f] != 3 &&
-         m->flex_stiffness[m->flex_stiffnessadr[f]] != 0;
+// return 1 if any flex needs implicit interp treatment
+static int flexInterp_has_active(const mjModel* m) {
+  for (int f=0; f < m->nflex; f++) {
+    if (m->flex_interp[f] && !m->flex_rigid[f] &&
+        m->flex_edgeequality[f] != 3 &&
+        m->flex_stiffness[m->flex_stiffnessadr[f]] != 0) {
+      return 1;
+    }
+  }
+  return 0;
 }
 
 
-// context for flex interp reduced banded factorization/solve
-typedef struct {
-  mjtNum* H;              // banded Cholesky-factored matrix (ndof x nband)
-  int* dof_indices;       // global DOF index for each local flex DOF
-  int ndof;               // number of flex DOFs
-  int nband;              // half-bandwidth + 1 (number of band columns)
-  int ncoupling;          // number of off-diagonal coupling terms
-  mjtNum* coupling_val;   // coupling coefficient values
-  int* coupling_row;      // local flex row index for each coupling term
-  int* coupling_col;      // global DOF column index for each coupling term
-} FlexInterpContext;
-
-
-// collect flex DOFs for one flex, marking seen_dof and incrementing count
-static void flexInterp_collect(const mjModel* m, int f,
-                               int* chain_dofs, int* seen_dof, int* count) {
-  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 chain_nnz;
-    if (m->body_dofnum[b] == 0) {
-      // pinned node: use bodyChain to get parent DOFs
-      chain_nnz = mj_bodyChain(m, b, chain_dofs);
-    } else {
-      // regular flex node: use body's own DOFs only
-      chain_nnz = m->body_dofnum[b];
-      for (int j=0; j < chain_nnz; j++) {
-        chain_dofs[j] = m->body_dofadr[b] + j;
-      }
-    }
-    for (int i=0; i < chain_nnz; i++) {
-      int dof = chain_dofs[i];
-      if (!seen_dof[dof]) {
-        seen_dof[dof] = 1;
-        (*count)++;
-      }
-    }
-  }
-}
-
-
-// build and factor the reduced banded matrix for flex interp DOFs
-//   mark/free stack handled by caller
-static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) {
-  FlexInterpContext ctx = {0};
-
-  int* chain_dofs = mjSTACKALLOC(d, nv, int);
-  int* seen_dof = mjSTACKALLOC(d, nv, int);
-  mju_fillInt(seen_dof, 0, nv);
-
-  // count flex DOFs
-  int ndof = 0;
-  for (int f=0; f < m->nflex; f++) {
-    if (flexInterp_active(m, f)) {
-      flexInterp_collect(m, f, chain_dofs, seen_dof, &ndof);
-    }
-  }
-  if (ndof == 0) {
-    return ctx;
-  }
-
-  // allocate and build global-to-local mapping
-  int* dof_indices = mjSTACKALLOC(d, ndof, int);
-  int* global2local = mjSTACKALLOC(d, nv, int);
-  mju_fillInt(global2local, -1, nv);
-
-  // collect unique DOFs in order
-  int cnt = 0;
-  mju_fillInt(seen_dof, 0, nv);
-  for (int f=0; f < m->nflex; f++) {
-    if (flexInterp_active(m, 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 chain_nnz;
-        if (m->body_dofnum[b] == 0) {
-          // pinned node: use bodyChain to get parent DOFs
-          chain_nnz = mj_bodyChain(m, b, chain_dofs);
-        } else {
-          // regular flex node: use body's own DOFs only
-          chain_nnz = m->body_dofnum[b];
-          for (int j=0; j < chain_nnz; j++) {
-            chain_dofs[j] = m->body_dofadr[b] + j;
-          }
-        }
-        for (int i=0; i < chain_nnz; i++) {
-          int dof = chain_dofs[i];
-          if (!seen_dof[dof]) {
-            seen_dof[dof] = 1;
-            dof_indices[cnt] = dof;
-            global2local[dof] = cnt;
-            cnt++;
-          }
-        }
-      }
-    }
-  }
-
-  // select sparse matrix format based on integrator
-  int implicit = (m->opt.integrator == mjINT_IMPLICIT);
-  const int* rownnz = implicit ? m->D_rownnz : m->M_rownnz;
-  const int* rowadr = implicit ? m->D_rowadr : m->M_rowadr;
-  const int* colind = implicit ? m->D_colind : m->M_colind;
-  const mjtNum* source = implicit ? d->qLU : d->qH;
-
-  // get precomputed bandwidth
-  int bandwidth = 0;
-  for (int f=0; f < m->nflex; f++) {
-    if (flexInterp_active(m, f)) {
-      if (m->flex_bandwidth[f] > bandwidth) {
-        bandwidth = m->flex_bandwidth[f];
-      }
-    }
-  }
-
-  // compute ncoupling from sparse matrix entries
-  int ncoupling = 0;
-  for (int i=0; i < ndof; i++) {
-    int row = dof_indices[i];
-    int start = rowadr[row];
-    int end = start + rownnz[row];
-    for (int k=start; k < end; k++) {
-      int local_j = global2local[colind[k]];
-      if (local_j < 0) {
-        ncoupling++;
-      }
-    }
-  }
-
-  // nband = bandwidth + 1 (includes diagonal)
-  int nband = bandwidth + 1;
-
-  // cap nband at ndof (dense fallback for small systems)
-  if (nband > ndof) nband = ndof;
-
-  // allocate coupling storage
-  mjtNum* coupling_val = NULL;
-  int* coupling_row = NULL;
-  int* coupling_col = NULL;
-  if (ncoupling > 0) {
-    coupling_val = mjSTACKALLOC(d, ncoupling, mjtNum);
-    coupling_row = mjSTACKALLOC(d, ncoupling, int);
-    coupling_col = mjSTACKALLOC(d, ncoupling, int);
-  }
-
-  // build H_flex (banded) from qLU (implicit) or qH (implicitfast)
-  mjtNum* H = mjSTACKALLOC(d, ndof*nband, mjtNum);
-  mju_zero(H, ndof*nband);
-
-  int coup_cnt = 0;
-  for (int i=0; i < ndof; i++) {
-    int row = 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) {
-        // store lower triangle only: row i, col local_j, where i >= local_j
-        if (i >= local_j) {
-          H[i*nband + nband-1-(i-local_j)] = source[k];
-        } else {
-          // upper triangle entry: store symmetrically in lower triangle
-          H[local_j*nband + nband-1-(local_j-i)] = source[k];
-        }
-      } else if (coup_cnt < ncoupling) {
-        coupling_val[coup_cnt] = source[k];
-        coupling_row[coup_cnt] = i;
-        coupling_col[coup_cnt] = col;
-        coup_cnt++;
-      }
-    }
-  }
-
-  // add flex stiffness in banded format and factorize
-  mjd_flexInterp_addH(m, d, H, dof_indices, ndof, nband, m->opt.timestep);
-  mju_cholFactorBand(H, ndof, nband, 0, 0, 0);
-
-  // store results in context
-  ctx.H = H;
-  ctx.dof_indices = dof_indices;
-  ctx.ndof = ndof;
-  ctx.nband = nband;
-  ctx.ncoupling = ncoupling;
-  ctx.coupling_val = coupling_val;
-  ctx.coupling_row = coupling_row;
-  ctx.coupling_col = coupling_col;
-  return ctx;
-}
-
-
-// solve the reduced banded system for flex interp DOFs, overwrite qacc
-static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContext* ctx,
-                             mjtNum* qacc, const mjtNum* qfrc, int nv) {
-  int ndof = ctx->ndof;
-  mjtNum* qfrc_flex = mjSTACKALLOC(d, ndof, mjtNum);
-  mjtNum* res = mjSTACKALLOC(d, nv, mjtNum);
-
+// preconditioned CG solve for implicit flex interp
+//   solves (M - h*qDeriv - (h^2+h*d)*K) * qacc = qfrc - h*K*qvel
+//   where K is the flex stiffness, using the already-factored standard system
+//   (M - h*qDeriv) as a preconditioner
+static void flexInterp_cgsolve(const mjModel* m, mjData* d,
+                                mjtNum* qacc, const mjtNum* qfrc, int nv) {
   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;
+  int implicit = (m->opt.integrator == mjINT_IMPLICIT);
 
-  // velocity correction: -h * K * v
-  mju_zero(res, nv);
-  mjd_flexInterp_mulKD(m, d, res, d->qvel, h);
+  mj_markStack(d);
 
-  for (int i=0; i < ndof; i++) {
-    int global_dof = ctx->dof_indices[i];
-    qfrc_flex[i] = qfrc[global_dof] + res[global_dof] * factor;
+  // allocate CG work vectors
+  mjtNum* rhs = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* r = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* z = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* p = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* Ap = mjSTACKALLOC(d, nv, mjtNum);
+  mjtNum* temp = mjSTACKALLOC(d, nv, mjtNum);
+
+  // build RHS: rhs = qfrc - h*K*qvel  (velocity correction from flex stiffness)
+  mju_copy(rhs, qfrc, nv);
+  mju_zero(temp, nv);
+  mjd_flexInterp_mulK(m, d, temp, d->qvel, h);  // temp = h*K*v (stiffness only)
+  mju_addToScl(rhs, temp, -1.0, nv);  // rhs -= h*K*v
+
+  // --- helper lambda-style inline: compute Ap = A*x ---
+  // A*x = (M - h*qDeriv)*x - (h^2+h*d)*K*x
+  #define FLEX_CG_MATVEC(Ap_out, x_in)                                           \
+    mju_mulMatVecSparse(Ap_out, d->qDeriv, x_in, nv, m->D_rownnz, m->D_rowadr,   \
+                        m->D_colind, NULL);                                      \
+    mju_zero(temp, nv);                                                          \
+    mju_mulSymVecSparse(temp, d->M, x_in, nv, m->M_rownnz, m->M_rowadr,          \
+                        m->M_colind);                                            \
+    mju_addScl(Ap_out, temp, Ap_out, -h, nv);                                    \
+    mju_zero(temp, nv);                                                          \
+    mjd_flexInterp_mulKD(m, d, temp, x_in, h);                                   \
+    mju_addToScl(Ap_out, temp, -1.0, nv)
+
+  // --- helper: preconditioner solve z = (M - h*qDeriv)^{-1} * r ---
+  #define FLEX_CG_PRECOND(z_out, r_in)                                           \
+    if (implicit) {                                                              \
+      mju_solveLUSparse(z_out, d->qLU, r_in, nv, m->D_rownnz, m->D_rowadr,       \
+                        m->D_diag, m->D_colind, NULL);                           \
+    } else {                                                                     \
+      mju_copy(z_out, r_in, nv);                                                 \
+      mj_solveLD(z_out, d->qH, d->qHDiagInv, nv, 1, m->M_rownnz, m->M_rowadr,    \
+                 m->M_colind, NULL);                                             \
+    }
+
+  // initial residual: r = rhs - A*qacc
+  FLEX_CG_MATVEC(Ap, qacc);
+  mju_sub(r, rhs, Ap, nv);
+
+  // check if already converged
+  mjtNum rnorm = mju_dot(r, r, nv);
+  mjtNum tol = 1e-10 * mju_dot(rhs, rhs, nv);
+  if (rnorm < tol || rnorm < mjMINVAL) {
+    mj_freeStack(d);
+    return;
   }
 
-  // coupling correction: qfrc_flex -= H_coupling * qacc_parent
-  for (int k=0; k < ctx->ncoupling; k++) {
-    qfrc_flex[ctx->coupling_row[k]] -= ctx->coupling_val[k] * qacc[ctx->coupling_col[k]];
+  // z = precond(r), p = z
+  FLEX_CG_PRECOND(z, r);
+  mju_copy(p, z, nv);
+  mjtNum rz = mju_dot(r, z, nv);
+
+  // CG iterations
+  int maxiter = 50;
+  for (int iter=0; iter < maxiter; iter++) {
+    FLEX_CG_MATVEC(Ap, p);
+
+    // alpha = rz / dot(p, Ap)
+    mjtNum pAp = mju_dot(p, Ap, nv);
+    if (mju_abs(pAp) < mjMINVAL) break;
+    mjtNum alpha = rz / pAp;
+
+    // qacc += alpha * p
+    mju_addToScl(qacc, p, alpha, nv);
+
+    // r -= alpha * Ap
+    mju_addToScl(r, Ap, -alpha, nv);
+
+    // check convergence
+    rnorm = mju_dot(r, r, nv);
+    if (rnorm < tol || rnorm < mjMINVAL) break;
+
+    // z = precond(r)
+    FLEX_CG_PRECOND(z, r);
+
+    // beta = rz_new / rz
+    mjtNum rz_new = mju_dot(r, z, nv);
+    mjtNum beta = rz_new / mju_max(mjMINVAL, rz);
+
+    // p = z + beta * p
+    mju_addScl(p, z, p, beta, nv);
+    rz = rz_new;
   }
 
-  // solve with banded Cholesky and scatter back
-  mju_cholSolveBand(qfrc_flex, ctx->H, qfrc_flex, ndof, ctx->nband, 0);
-  mju_scatter(qacc, qfrc_flex, ctx->dof_indices, ndof);
+  #undef FLEX_CG_MATVEC
+  #undef FLEX_CG_PRECOND
+
+  mj_freeStack(d);
 }
 
 
@@ -1972,16 +1857,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
   }
 
   // check for flex_interp that needs implicit treatment
-  int has_flex_interp = 0;
-  for (int f=0; f < m->nflex; f++) {
-    if (flexInterp_active(m, f)) {
-      has_flex_interp = 1;
-      break;
-    }
-  }
-
-  // flex interp context (populated during factorization)
-  FlexInterpContext flex = {0};
+  int has_flex_interp = !sleep_filter && flexInterp_has_active(m);
 
   // factorization
   if (!skipfactor) {
@@ -2011,11 +1887,6 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
       mjERROR("integrator must be implicit or implicitfast");
     }
 
-    // flex: reduced dense factorization
-    if (has_flex_interp && !sleep_filter) {
-      flex = flexInterp_factor(m, d, nv);
-    }
-
     // standard factorization (implicit / implicitfast)
     if (m->opt.integrator == mjINT_IMPLICIT) {
       int* scratch = mjSTACKALLOC(d, nv, int);
@@ -2039,9 +1910,9 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
     mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
   }
 
-  // flex: reduced dense solve
-  if (flex.H) {
-    flexInterp_solve(m, d, &flex, qacc, qfrc, nv);
+  // flex: CG correction for implicit flex stiffness
+  if (has_flex_interp) {
+    flexInterp_cgsolve(m, d, qacc, qfrc, m->nv);
   }
 
   // count and list joints of free bodies eligible for midpoint integration
diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c
index 81a0f155..d47c570e 100644
--- a/src/engine/engine_setconst.c
+++ b/src/engine/engine_setconst.c
@@ -638,135 +638,6 @@ static void makeFlexSparse(mjModel* m, mjData* d) {
   mj_freeStack(d);
 }
 
-// compute flex bandwidth for trilinear interpolation
-static void makeFlexBandwidth(mjModel* m, mjData* d) {
-  if (!m->nflex) {
-    return;
-  }
-
-  mj_markStack(d);
-  int* chain_dofs = mjSTACKALLOC(d, m->nv, int);
-  int* seen_dof = mjSTACKALLOC(d, m->nv, int);
-  int* dof_indices = mjSTACKALLOC(d, m->nv, int);
-  int* global2local = mjSTACKALLOC(d, m->nv, int);
-
-  mju_zeroInt(seen_dof, m->nv);
-  for (int i = 0; i < m->nv; i++) {
-    global2local[i] = -1;
-  }
-
-  int ndof = 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];
-        // only the body's own DOFs enter the reduced banded flex system;
-        // ancestor DOFs are solved by the global factorization and coupled
-        // via off-diagonal correction (see flexInterp_solve in engine_forward)
-        int chain_nnz;
-        if (m->body_dofnum[b] == 0) {
-          chain_nnz = mj_bodyChain(m, b, chain_dofs);
-        } else {
-          chain_nnz = m->body_dofnum[b];
-          for (int j = 0; j < chain_nnz; j++) {
-            chain_dofs[j] = m->body_dofadr[b] + j;
-          }
-        }
-        for (int i = 0; i < chain_nnz; i++) {
-          int dof = chain_dofs[i];
-          if (!seen_dof[dof]) {
-            seen_dof[dof] = 1;
-            dof_indices[ndof] = dof;
-            global2local[dof] = ndof++;
-          }
-        }
-      }
-    }
-  }
-
-  int bandwidth = 0;
-  if (ndof > 0) {
-    // check sparse matrix coupling (both D and M)
-    for (int integrator = 0; integrator < 2; integrator++) {
-      const int* rownnz = (integrator == 0) ? m->D_rownnz : m->M_rownnz;
-      const int* rowadr = (integrator == 0) ? m->D_rowadr : m->M_rowadr;
-      const int* colind = (integrator == 0) ? m->D_colind : m->M_colind;
-
-      // D arrays are only allocated for implicit integrators
-      if (!rownnz) continue;
-
-      for (int i = 0; i < ndof; i++) {
-        int row = dof_indices[i];
-        int start = rowadr[row];
-        int end = start + rownnz[row];
-        for (int k = start; k < end; k++) {
-          int local_j = global2local[colind[k]];
-          if (local_j >= 0) {
-            int diff = i - local_j;
-            if (diff < 0) diff = -diff;
-            if (diff > bandwidth) bandwidth = diff;
-          }
-        }
-      }
-    }
-
-    // check stiffness coupling
-    for (int f = 0; f < m->nflex; f++) {
-      if (!m->flex_interp[f]) continue;
-      int order = m->flex_interp[f];
-      order = order < 0 ? -order : order;
-      int nodeadr = m->flex_nodeadr[f];
-      int nodenum = m->flex_nodenum[f];
-      int cx = m->flex_cellnum[3*f+0];
-      int cy = m->flex_cellnum[3*f+1];
-      int cz = m->flex_cellnum[3*f+2];
-      int ny = cy * order + 1;
-      int nz = cz * order + 1;
-
-      for (int icx = 0; icx < cx; icx++) {
-        for (int icy = 0; icy < cy; icy++) {
-          for (int icz = 0; icz < cz; icz++) {
-            int min_local = ndof, max_local = -1;
-            for (int lx = 0; lx <= order; lx++) {
-              for (int ly = 0; ly <= order; ly++) {
-                for (int lz = 0; lz <= order; lz++) {
-                  int gx = icx * order + lx;
-                  int gy = icy * order + ly;
-                  int gz = icz * order + lz;
-                  int node_idx = gx * ny * nz + gy * nz + gz;  // non-negative by construction
-                  if (node_idx < nodenum) {
-                    int b = m->flex_nodebodyid[nodeadr + node_idx];
-                    int chain_nnz = mj_bodyChain(m, b, chain_dofs);
-                    for (int i = 0; i < chain_nnz; i++) {
-                      int dof = chain_dofs[i];
-                      int local = global2local[dof];
-                      if (local >= 0) {
-                        if (local < min_local) min_local = local;
-                        if (local > max_local) max_local = local;
-                      }
-                    }
-                  }
-                }
-              }
-            }
-            if (max_local >= 0 && max_local - min_local > bandwidth) {
-              bandwidth = max_local - min_local;
-            }
-          }
-        }
-      }
-    }
-  }
-
-  // store bandwidth for all flexes (global max)
-  for (int f = 0; f < m->nflex; f++) {
-    m->flex_bandwidth[f] = bandwidth;
-  }
-
-  mj_freeStack(d);
-}
 
 // align 2D flexes to the XY plane
 static void mj_alignFlex(mjModel* m, mjData* d) {
@@ -819,7 +690,7 @@ static void mj_alignFlex(mjModel* m, mjData* d) {
 static void set0(mjModel* m, mjData* d) {
   makeTendonSparse(m);
   makeFlexSparse(m, d);
-  makeFlexBandwidth(m, d);
+
   mj_alignFlex(m, d);
   int nv = m->nv;
   mjtNum A[36] = {0}, pos[3], quat[4];
diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc
index 0e5c1849..fb9d777f 100644
--- a/test/engine/engine_derivative_test.cc
+++ b/test/engine/engine_derivative_test.cc
@@ -1473,15 +1473,24 @@ void RotateFlexGrid(mjModel* model, mjData* data, const char* flex_name,
   }
 }
 
-// Helper: assemble flex stiffness into dense matrix via banded addH
-// This wraps the banded API and converts to dense for test verification.
-static void addH_dense(mjModel* m, mjData* d, mjtNum* H_dense,
-                       const int* dof_indices, int ndof, mjtNum h) {
-  // use full bandwidth (ndof) for exact dense equivalence
-  std::vector<mjtNum> H_band(ndof * ndof, 0);
-  mjd_flexInterp_addH(m, d, H_band.data(), dof_indices, ndof, ndof, h);
-  // convert banded to dense (lower triangle), then symmetrize
-  mju_band2Dense(H_dense, H_band.data(), ndof, ndof, 0, 1);
+// Helper: assemble flex stiffness into dense matrix via matrix-vector products.
+// Builds K column-by-column using mjd_flexInterp_mulKD.
+// Result is -(h^2 + h*damping) * J'KJ (negative sign matches the old addH
+// convention where stiffness is subtracted from the system matrix).
+static void mulKD_dense(mjModel* m, mjData* d, mjtNum* H_dense,
+                        int nv, mjtNum h) {
+  std::vector<mjtNum> e_i(nv, 0);
+  std::vector<mjtNum> col(nv, 0);
+  for (int i = 0; i < nv; i++) {
+    mju_zero(e_i.data(), nv);
+    mju_zero(col.data(), nv);
+    e_i[i] = 1.0;
+    mjd_flexInterp_mulKD(m, d, col.data(), e_i.data(), h);
+    // col = +(h^2 + h*damp)*K*e_i, negate to match addH convention (H -= K)
+    for (int j = 0; j < nv; j++) {
+      H_dense[j * nv + i] = -col[j];
+    }
+  }
 }
 
 // compare analytic and fin-diff d_qfrc_passive/d_qvel for flex interp
@@ -1525,18 +1534,16 @@ TEST_F(DerivativeTest, FlexInterpDerivatives) {
         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
+      // use mulKD to compute K * vec
+      // mulKD adds (h^2*K + h*D)*vec to res
+      // if we set h=1, damping=0, we get K*vec
       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
-      addH_dense(model, data, H.data(), dof_indices.data(), nv, 1.0);
+      // assemble K into H column-by-column
+      mulKD_dense(model, data, H.data(), nv, 1.0);
 
       // restore damping
       model->flex_damping[0] = save_damping;
@@ -1623,16 +1630,13 @@ TEST_F(DerivativeTest, FlexInterpDerivatives) {
       // 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
+      // compute expected flex damping using mulKD_dense
       // 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);
-      addH_dense(model, data, H1.data(), dof_indices.data(), nv, 1.0);
+      mulKD_dense(model, data, H1.data(), nv, 1.0);
 
       vector<mjtNum> H2(nv * nv, 0);
-      addH_dense(model, data, H2.data(), dof_indices.data(), nv, 0.5);
+      mulKD_dense(model, data, H2.data(), nv, 0.5);
 
       vector<mjtNum> D(nv * nv);
       for (int i = 0; i < nv * nv; i++) {
@@ -1700,13 +1704,11 @@ TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) {
   mj_forward(model, data);
 
   // 1. Compute Analytic Jacobian (Approximate)
-  // We use mjd_flexInterp_addH to get K_approx
+  // We use mulKD_dense 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
-  addH_dense(model, data, H_approx.data(), dof_indices.data(), nv, 1.0);
+  // h=1, damping=0 => gives K
+  mulKD_dense(model, data, H_approx.data(), nv, 1.0);
 
   // 2. Compute Finite Difference Jacobian (Ground Truth)
   // qfrc_passive = -dV/dq
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index a655529c..d546c345 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -1201,7 +1201,6 @@ public unsafe struct mjModel_ {
   public int* flex_matid;
   public int* flex_group;
   public int* flex_interp;
-  public int* flex_bandwidth;
   public int* flex_cellnum;
   public int* flex_nodeadr;
   public int* flex_nodenum;
diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc
index f0c11146..e78ece4f 100644
--- a/wasm/codegen/generated/bindings.cc
+++ b/wasm/codegen/generated/bindings.cc
@@ -4655,9 +4655,6 @@ struct MjModel {
   emscripten::val flex_interp() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_interp));
   }
-  emscripten::val flex_bandwidth() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_bandwidth));
-  }
   emscripten::val flex_cellnum() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nflex * 3, ptr_->flex_cellnum));
   }
@@ -11854,7 +11851,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("eq_type", &MjModel::eq_type)
     .property("exclude_signature", &MjModel::exclude_signature)
     .property("flex_activelayers", &MjModel::flex_activelayers)
-    .property("flex_bandwidth", &MjModel::flex_bandwidth)
     .property("flex_bending", &MjModel::flex_bending)
     .property("flex_bendingadr", &MjModel::flex_bendingadr)
     .property("flex_bvhadr", &MjModel::flex_bvhadr)