diff --git a/doc/includes/references.h b/doc/includes/references.h index 614a7516..0ee5a22b 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -1315,6 +1315,7 @@ 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 9b39cbb6..fba680f9 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -978,6 +978,7 @@ 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 87a95a31..4efaf94b 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -454,6 +454,7 @@ 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/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index 14a44f32..9bd45295 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -2668,6 +2668,14 @@ 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 1eb4df9d..8a7eabc7 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -872,12 +872,12 @@ typedef enum { // 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) +// res: output vector (VEC) or banded H matrix (ADDH) // vec: input vector for VEC operation, NULL otherwise -// dof_indices, ndof: DOF mapping for ADDH, ignored otherwise +// dof_indices, ndof, nband: DOF mapping and band width 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) { + const int* dof_indices, int ndof, int nband) { int nv = m->nv; // build global2local map for ADDH @@ -1021,7 +1021,7 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op, J_val, K_rot_cell, dim_c); } else if (op == mjFLEXOP_ADDH) { mj_markStack(d); - // H -= J_cell^T * K_rot_cell * J_cell + // H -= J_cell^T * K_rot_cell * J_cell (banded format) mjtNum* J_reduced = mjSTACKALLOC(d, dim_c*ndof, mjtNum); mju_zero(J_reduced, dim_c*ndof); @@ -1041,14 +1041,14 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op, mjtNum* KJ = mjSTACKALLOC(d, dim_c*ndof, mjtNum); mju_mulMatMat(KJ, K_rot_cell, J_reduced, dim_c, dim_c, ndof); - // H[i,j] -= J_reduced[k,i] * KJ[k,j] + // H[i,j] -= J_reduced[k,i] * KJ[k,j], store lower triangle in banded format for (int i = 0; i < ndof; i++) { - for (int j = 0; j < ndof; j++) { + for (int j = mjMAX(0, i-nband+1); j <= i; j++) { mjtNum val = 0; for (int dim_idx = 0; dim_idx < dim_c; dim_idx++) { val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j]; } - res[i*ndof + j] -= val; + res[i*nband + nband-1-(i-j)] -= val; } } mj_freeStack(d); @@ -1072,15 +1072,16 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op, // 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); + mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, h * h, h, NULL, 0, 0); } -// add (h^2 + h*damping) * J'*K*J to dense matrix H, for all interpolated flexes -// H: dense ndof x ndof matrix +// 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, mjtNum h) { - mjd_flexInterp_kernel(m, d, mjFLEXOP_ADDH, H, NULL, h * h, h, dof_indices, ndof); +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); } diff --git a/src/engine/engine_derivative.h b/src/engine/engine_derivative.h index af2aef4c..65a136f5 100644 --- a/src/engine/engine_derivative.h +++ b/src/engine/engine_derivative.h @@ -49,7 +49,7 @@ MJAPI void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const // 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); +MJAPI void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, int nband, mjtNum h); #ifdef __cplusplus diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index b0d15b16..dbd3a74e 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -1350,11 +1350,12 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) { } -// context for flex interp reduced dense factorization/solve +// context for flex interp reduced banded factorization/solve typedef struct { - mjtNum* H; // dense Cholesky-factored matrix (ndof x ndof) + 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 @@ -1391,7 +1392,7 @@ static void flexInterp_collect(const mjModel* m, int f, } -// build and factor the reduced dense matrix for flex interp DOFs +// 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}; @@ -1456,19 +1457,36 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) const int* colind = implicit ? m->D_colind : m->M_colind; const mjtNum* source = implicit ? d->qLU : d->qH; - // count coupling terms (off-diagonal: flex row, non-flex col) + // get precomputed bandwidth + int bandwidth = 0; + for (int f=0; f < m->nflex; f++) { + if (m->flex_interp[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++) { - if (global2local[colind[k]] < 0) { + 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; @@ -1479,9 +1497,9 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) coupling_col = mjSTACKALLOC(d, ncoupling, int); } - // build H_flex (dense) from qLU (implicit) or qH (implicitfast) - mjtNum* H = mjSTACKALLOC(d, ndof*ndof, mjtNum); - mju_zero(H, ndof*ndof); + // 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++) { @@ -1492,7 +1510,13 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) int col = colind[k]; int local_j = global2local[col]; if (local_j >= 0) { - H[i*ndof+local_j] = source[k]; + // 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; @@ -1502,14 +1526,15 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) } } - // add flex stiffness and factorize - mjd_flexInterp_addH(m, d, H, dof_indices, ndof, m->opt.timestep); - mju_cholFactor(H, ndof, mjMINVAL); + // 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; @@ -1518,7 +1543,7 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) } -// solve the reduced dense system for flex interp DOFs, overwrite qacc +// 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; @@ -1544,8 +1569,8 @@ static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContex qfrc_flex[ctx->coupling_row[k]] -= ctx->coupling_val[k] * qacc[ctx->coupling_col[k]]; } - // solve and scatter back - mju_cholSolve(qfrc_flex, ctx->H, qfrc_flex, ndof); + // 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); } diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index e0ffb480..24ef9955 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -636,6 +636,135 @@ 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]; + 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) { for (int f = 0; f < m->nflex; f++) { @@ -687,6 +816,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/src/user/user_mesh.cc b/src/user/user_mesh.cc index 51a940a4..244a48f2 100644 --- a/src/user/user_mesh.cc +++ b/src/user/user_mesh.cc @@ -4095,6 +4095,9 @@ void mjCFlex::Compile(const mjVFS* vfs) { // check node compatibility with count and dof if (spec.order > 0) { + if (spec.cellcount[0] == 0 || spec.cellcount[1] == 0 || spec.cellcount[2] == 0) { + throw mjCError(this, "cellcount cannot be 0 in any dimension when interpolation order > 0"); + } int expected_nodes = (spec.cellcount[0] * spec.order + 1) * (spec.cellcount[1] * spec.order + 1) * (spec.cellcount[2] * spec.order + 1); diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index e5978c9b..e0b693d3 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -1473,6 +1473,17 @@ 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 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); +} + // 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) { @@ -1525,7 +1536,7 @@ TEST_F(DerivativeTest, FlexInterpDerivatives) { 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); + addH_dense(model, data, H.data(), dof_indices.data(), nv, 1.0); // restore damping model->flex_damping[0] = save_damping; @@ -1618,10 +1629,10 @@ TEST_F(DerivativeTest, FlexInterpDerivatives) { 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); + addH_dense(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); + addH_dense(model, data, H2.data(), dof_indices.data(), nv, 0.5); vector D(nv * nv); for (int i = 0; i < nv * nv; i++) { @@ -1695,8 +1706,7 @@ TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) { 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); + addH_dense(model, data, H_approx.data(), dof_indices.data(), nv, 1.0); // 2. Compute Finite Difference Jacobian (Ground Truth) // qfrc_passive = -dV/dq diff --git a/test/user/user_flex_test.cc b/test/user/user_flex_test.cc index 478b23e0..5482568d 100644 --- a/test/user/user_flex_test.cc +++ b/test/user/user_flex_test.cc @@ -79,6 +79,27 @@ TEST_F(UserFlexTest, CountTooSmall) { EXPECT_THAT(error.data(), HasSubstr("Count too small")); } +TEST_F(UserFlexTest, CellnumZeroInterpolated) { + static constexpr char xml[] = R"( + + + + + + + + + + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(m, IsNull()); + EXPECT_THAT(error.data(), HasSubstr("cellcount cannot be 0")); +} + TEST_F(UserFlexTest, SpacingGreaterThanGeometry) { static constexpr char xml[] = R"( diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index d7d4c47d..1673c4d0 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -1200,6 +1200,7 @@ 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 80fb63d5..185f5bb7 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -4649,6 +4649,9 @@ 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)); } @@ -11815,6 +11818,7 @@ 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_bvhadr", &MjModel::flex_bvhadr) .property("flex_bvhnum", &MjModel::flex_bvhnum)