From 04042d8bf3b0f7138a291809326d0ddeffc7c508 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sun, 17 May 2026 11:21:07 -0700 Subject: [PATCH] Move square root of Delassus matrix from stack to arena PiperOrigin-RevId: 916852244 Change-Id: I4379553f808d0b238a1421606f806ea2d0a33d7c --- doc/includes/references.h | 7 +- include/mujoco/mjdata.h | 7 +- include/mujoco/mjxmacro.h | 5 + python/mujoco/introspect/structs.py | 39 ++- python/mujoco/structs_wrappers.cc | 2 + src/engine/engine_core_constraint.c | 378 ++++++++++++++------------- src/engine/engine_io.c | 1 + src/engine/engine_print.c | 11 +- unity/Runtime/Bindings/MjBindings.cs | 5 + wasm/codegen/generated/bindings.cc | 23 ++ wasm/codegen/generators/constants.py | 4 + 11 files changed, 294 insertions(+), 188 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 3034e6c9..b1c1f0df 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -191,6 +191,7 @@ struct mjData_ { int nl; // number of limit constraints int nefc; // number of constraints int nJ; // number of non-zeros in constraint Jacobian + int nY; // number of non-zeros in constraint inverse inertia square root int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands int nidof; // number of dofs in all islands @@ -445,8 +446,12 @@ struct mjData_ { mjtNum* iefc_R; // inverse constraint mass (nefc x 1) // computed by mj_projectConstraint (PGS solver) + int* efc_Y_rownnz; // number of non-zeros in Y row (nefc x 1) + int* efc_Y_rowadr; // row start address in Y colind array (nefc x 1) + int* efc_Y_colind; // column indices in sparse Y (nY x 1) + mjtNum* efc_Y; // whitened Jacobian Y = J*M^(-1/2) (nY x 1) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) - int* efc_AR_rowadr; // row start address in colind array (nefc x 1) + int* efc_AR_rowadr; // row start address in AR colind array (nefc x 1) int* efc_AR_colind; // column indices in sparse AR (nA x 1) mjtNum* efc_AR; // J*inv(M)*J' + R (nA x 1) diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index d092a27d..3709896a 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -225,6 +225,7 @@ struct mjData_ { int nl; // number of limit constraints int nefc; // number of constraints int nJ; // number of non-zeros in constraint Jacobian + int nY; // number of non-zeros in constraint inverse inertia square root int nA; // number of non-zeros in constraint inverse inertia matrix int nisland; // number of detected constraint islands int nidof; // number of dofs in all islands @@ -479,8 +480,12 @@ struct mjData_ { mjtNum* iefc_R; // inverse constraint mass (nefc x 1) // computed by mj_projectConstraint (PGS solver) + int* efc_Y_rownnz; // number of non-zeros in Y row (nefc x 1) + int* efc_Y_rowadr; // row start address in Y colind array (nefc x 1) + int* efc_Y_colind; // column indices in sparse Y (nY x 1) + mjtNum* efc_Y; // whitened Jacobian Y = J*M^(-1/2) (nY x 1) int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1) - int* efc_AR_rowadr; // row start address in colind array (nefc x 1) + int* efc_AR_rowadr; // row start address in AR colind array (nefc x 1) int* efc_AR_colind; // column indices in sparse AR (nA x 1) mjtNum* efc_AR; // J*inv(M)*J' + R (nA x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index c11ff844..98b5e75b 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -946,6 +946,10 @@ // array fields of mjData that are used in the dual problem #define MJDATA_ARENA_POINTERS_DUAL \ + XNV( int, efc_Y_rownnz, MJ_D(nefc), 1 ) \ + XNV( int, efc_Y_rowadr, MJ_D(nefc), 1 ) \ + XNV( int, efc_Y_colind, MJ_D(nY), 1 ) \ + XNV( mjtNum, efc_Y, MJ_D(nY), 1 ) \ XNV( int, efc_AR_rownnz, MJ_D(nefc), 1 ) \ XNV( int, efc_AR_rowadr, MJ_D(nefc), 1 ) \ XNV( int, efc_AR_colind, MJ_D(nA), 1 ) \ @@ -1020,6 +1024,7 @@ X( int, nl ) \ X( int, nefc ) \ X( int, nJ ) \ + X( int, nY ) \ X( int, nA ) \ X( int, nisland ) \ X( int, nidof ) \ diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index a9ee204a..0b924bf7 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -5492,6 +5492,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of non-zeros in constraint Jacobian', ), + StructFieldDecl( + name='nY', + type=ValueType(name='int'), + doc='number of non-zeros in constraint inverse inertia square root', # pylint: disable=line-too-long + ), StructFieldDecl( name='nA', type=ValueType(name='int'), @@ -6726,6 +6731,38 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='inverse constraint mass', array_extent=('nefc',), ), + StructFieldDecl( + name='efc_Y_rownnz', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='number of non-zeros in Y row', + array_extent=('nefc',), + ), + StructFieldDecl( + name='efc_Y_rowadr', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='row start address in Y colind array', + array_extent=('nefc',), + ), + StructFieldDecl( + name='efc_Y_colind', + type=PointerType( + inner_type=ValueType(name='int'), + ), + doc='column indices in sparse Y', + array_extent=('nY',), + ), + StructFieldDecl( + name='efc_Y', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + doc='whitened Jacobian Y = J*M^(-1/2)', + array_extent=('nY',), + ), StructFieldDecl( name='efc_AR_rownnz', type=PointerType( @@ -6739,7 +6776,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='int'), ), - doc='row start address in colind array', + doc='row start address in AR colind array', array_extent=('nefc',), ), StructFieldDecl( diff --git a/python/mujoco/structs_wrappers.cc b/python/mujoco/structs_wrappers.cc index 0f46d811..1a35c0dc 100644 --- a/python/mujoco/structs_wrappers.cc +++ b/python/mujoco/structs_wrappers.cc @@ -812,6 +812,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const { X(ne); X(nf); X(nJ); + X(nY); X(nA); X(nefc); X(nisland); @@ -891,6 +892,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) { X(ne); X(nf); X(nJ); + X(nY); X(nA); X(nefc); X(nisland); diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 2f94818f..428c6fdc 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -2591,12 +2591,137 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) { } +// pre-count Y_rownnz, Y_rowadr, return total nonzeros nY +// Y has the sparsity of J * inv(L'), where L is the Cholesky factor of M +static int computeY_precount(int* Y_rownnz, int* Y_rowadr, int nefc, int nv, + const int* J_rownnz, const int* J_rowadr, const int* J_colind, + const int* M_rownnz, const int* M_rowadr, const int* M_colind, + int* marker) { + mju_fillInt(marker, -1, nv); + + Y_rowadr[0] = 0; + for (int r=0; r < nefc; r++) { + int nnz = 0; // nonzeros in row r of Y + + // traverse row r of J in reverse, count unique nonzeros + int start = J_rowadr[r]; + int end = start + J_rownnz[r]; + for (int i=end-1; i >= start; i--) { + int j = J_colind[i]; + + // if dof j is marked, it was already counted by a child dof: skip it + if (marker[j] == r) { + continue; + } + + // traverse row j of M, marking new unique nonzeros + int nnzM = M_rownnz[j]; + int adrM = M_rowadr[j]; + for (int k=0; k < nnzM; k++) { + int c = M_colind[adrM + k]; + if (marker[c] != r) { + marker[c] = r; + nnz++; + } + } + } + + // update rownnz and rowadr + Y_rownnz[r] = nnz; + if (r < nefc - 1) { + Y_rowadr[r+1] = Y_rowadr[r] + nnz; + } + } + + // total non-zeros in Y + return Y_rowadr[nefc-1] + Y_rownnz[nefc-1]; +} + + +// fill Y column indices and values from J, chaining up the kinematic tree +static void computeY_fill(mjtNum* Y, int* Y_colind, + const int* Y_rownnz, const int* Y_rowadr, int nefc, + const mjtNum* J, const int* J_rownnz, const int* J_rowadr, + const int* J_colind, const int* dof_parentid) { + for (int r=0; r < nefc; r++) { + // init row + int end = Y_rowadr[r] + Y_rownnz[r]; + int adrJ = J_rowadr[r]; + int remainJ = J_rownnz[r]; + int nnzY = 0; + + // complete chain in reverse + while (1) { + // get previous dof in src and dst + int prev_src = (remainJ > 0 ? J_colind[adrJ + remainJ - 1] : -1); + int prev_dst = (nnzY > 0 ? dof_parentid[Y_colind[end - nnzY]] : -1); + + // both finished: break + if (prev_src < 0 && prev_dst < 0) { + break; + } + + // add src + else if (prev_src >= prev_dst) { + nnzY++; + remainJ--; + Y_colind[end - nnzY] = prev_src; + Y[end - nnzY] = J[adrJ + remainJ]; + } + + // add dst + else { + nnzY++; + Y_colind[end - nnzY] = prev_dst; + Y[end - nnzY] = 0; + } + } + + // compare with Y_rownnz: SHOULD NOT OCCUR + if (nnzY != Y_rownnz[r]) { + mjERROR("pre and post-count of Y_rownnz are not equal on row %d", r); + } + } +} + + +// in-place sparse back-substitution: Y <- Y * M^{-1/2} +static void computeY_backsub(mjtNum* Y, const int* Y_rownnz, const int* Y_rowadr, + const int* Y_colind, int nefc, + const mjtNum* qLD, const int* M_rownnz, const int* M_rowadr, + const int* M_colind, const mjtNum* sqrtInvD) { + for (int r=0; r < nefc; r++) { + int nnzY = Y_rownnz[r]; + int adrY = Y_rowadr[r]; + + // Y(r,:) <- inv(L') * Y(r,:), exploit sparsity of input vector + for (int i=adrY + nnzY-1; i >= adrY; i--) { + mjtNum val = Y[i]; + if (val == 0) { + continue; + } + int j = Y_colind[i]; + int adrM = M_rowadr[j]; + mju_addToSclSparseInc(Y + adrY, qLD + adrM, + nnzY, Y_colind + adrY, + M_rownnz[j]-1, M_colind + adrM, -val); + } + + // Y(r,:) <- sqrt(inv(D)) * Y(r,:) + for (int i=adrY; i < adrY + nnzY; i++) { + int j = Y_colind[i]; + Y[i] *= sqrtInvD[j]; + } + } +} + + //---------------------------- top-level API for constraint construction --------------------------- // driver: call all functions above void mj_makeConstraint(const mjModel* m, mjData* d) { // clear sizes - d->ne = d->nf = d->nl = d->nefc = d->nJ = d->nA = 0; + d->ne = d->nf = d->nl = d->nefc = d->nJ = d->nA = d->nY = 0; // disabled or Jacobian not allocated: return if (mjDISABLED(mjDSBL_CONSTRAINT)) { @@ -2705,177 +2830,60 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { sqrtInvD[i] = 1 / mju_sqrt(d->qLD[diag]); } - // sparse + // sparse Y = backsubM2(J')' and its transpose if (mj_isSparse(m)) { - // compute B = backsubM2(J')' and its transpose - - - // === pre-count B_rownnz, B_rowadr, nB (total nonzeros) - - // allocate B rownnz and rowadr - int* B_rownnz = mjSTACKALLOC(d, nefc, int); - int* B_rowadr = mjSTACKALLOC(d, nefc, int); + // arena-allocate Y rownnz and rowadr + d->efc_Y_rownnz = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int)); + d->efc_Y_rowadr = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int)); + if (!d->efc_Y_rownnz || !d->efc_Y_rowadr) { + mj_warning(d, mjWARN_CNSTRFULL, d->narena); + mj_clearEfc(d); + d->parena = d->ncon * sizeof(mjContact); + mj_freeStack(d); + return; + } // markers for merged dofs, initialized to -1 int* marker = mjSTACKALLOC(d, nv, int); - mju_fillInt(marker, -1, nv); - B_rowadr[0] = 0; - for (int r=0; r < nefc; r++) { - // supernode: same sparsity as previous row - if (r > 0 && d->efc_J_rowsuper[r-1] > 0) { - B_rownnz[r] = B_rownnz[r-1]; - } + // pre-count Y_rownnz, Y_rowadr, nY (total nonzeros) + d->nY = computeY_precount(d->efc_Y_rownnz, d->efc_Y_rowadr, nefc, nv, + d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, + m->M_rownnz, m->M_rowadr, m->M_colind, marker); - // first row in supernode block: full chain traversal - else { - int nnz = 0; - - // traverse row r of J in reverse, count unique nonzeros - int start = d->efc_J_rowadr[r]; - int end = start + d->efc_J_rownnz[r]; - for (int i=end-1; i >= start; i--) { - int j = d->efc_J_colind[i]; - - // if dof j is marked, it was already counted by a child dof: skip it - if (marker[j] == r) { - continue; - } - - // traverse row j of M, marking new unique nonzeros - int nnzM = m->M_rownnz[j]; - int adrM = m->M_rowadr[j]; - for (int k=0; k < nnzM; k++) { - int c = m->M_colind[adrM + k]; - if (marker[c] != r) { - marker[c] = r; - nnz++; - } - } - } - B_rownnz[r] = nnz; - } - - // update rowadr - if (r < nefc - 1) { - B_rowadr[r+1] = B_rowadr[r] + B_rownnz[r]; - } + // arena-allocate values and column indices + d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum)); + d->efc_Y_colind = mj_arenaAllocByte(d, sizeof(int) * d->nY, _Alignof(int)); + if (!d->efc_Y || !d->efc_Y_colind) { + mj_warning(d, mjWARN_CNSTRFULL, d->narena); + mj_clearEfc(d); + d->parena = d->ncon * sizeof(mjContact); + mj_freeStack(d); + return; } - // total non-zeros in B - int nB = B_rowadr[nefc-1] + B_rownnz[nefc-1]; + // fill in Y column indices, copy values from J + computeY_fill(d->efc_Y, d->efc_Y_colind, d->efc_Y_rownnz, d->efc_Y_rowadr, nefc, + d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, + m->dof_parentid); - // === fill in B column indices, copy values from J + // in-place sparse back-substitution: Y <- Y * M^-1/2 + computeY_backsub(d->efc_Y, d->efc_Y_rownnz, d->efc_Y_rowadr, + d->efc_Y_colind, nefc, + d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind, sqrtInvD); - // allocate values and column indices - mjtNum* B = mjSTACKALLOC(d, nB, mjtNum); - int* B_colind = mjSTACKALLOC(d, nB, int); + // Y supernodes are identical to J supernodes + const int* Y_rowsuper = d->efc_J_rowsuper; - for (int r=0; r < nefc; r++) { - // supernode: copy column indices, only update values from J - if (r > 0 && d->efc_J_rowsuper[r-1] > 0) { - int prevAdr = B_rowadr[r-1]; - int adrB = B_rowadr[r]; - int nnzB = B_rownnz[r]; - mju_copyInt(B_colind + adrB, B_colind + prevAdr, nnzB); - mju_zero(B + adrB, nnzB); - - // copy J values into correct positions - int adrJ = d->efc_J_rowadr[r]; - int jnnz = d->efc_J_rownnz[r]; - int bi = 0, ji = 0; - while (ji < jnnz && bi < nnzB) { - if (B_colind[adrB+bi] == d->efc_J_colind[adrJ+ji]) { - B[adrB+bi] = d->efc_J[adrJ+ji]; - bi++; - ji++; - } else { - bi++; - } - } - } - - // first row in supernode block: full chain completion - else { - int end = B_rowadr[r] + B_rownnz[r]; - int adrJ = d->efc_J_rowadr[r]; - int remainJ = d->efc_J_rownnz[r]; - int nnzB = 0; - - // complete chain in reverse - while (1) { - // get previous dof in src and dst - int prev_src = (remainJ > 0 ? d->efc_J_colind[adrJ + remainJ - 1] : -1); - int prev_dst = (nnzB > 0 ? m->dof_parentid[B_colind[end - nnzB]] : -1); - - // both finished: break - if (prev_src < 0 && prev_dst < 0) { - break; - } - - // add src - else if (prev_src >= prev_dst) { - nnzB++; - remainJ--; - B_colind[end - nnzB] = prev_src; - B[end - nnzB] = d->efc_J[adrJ + remainJ]; - } - - // add dst - else { - nnzB++; - B_colind[end - nnzB] = prev_dst; - B[end - nnzB] = 0; - } - } - - // compare with B_rownnz: SHOULD NOT OCCUR - if (nnzB != B_rownnz[r]) { - mjERROR("pre and post-count of B_rownnz are not equal on row %d", r); - } - } - } - - - // === in-place sparse back-substitution: B <- B * M^-1/2 - - // sparse backsubM2 (half of LD back-substitution) - for (int r=0; r < nefc; r++) { - int nnzB = B_rownnz[r]; - int adrB = B_rowadr[r]; - - // B(r,:) <- inv(L') * B(r,:), exploit sparsity of input vector - for (int i=adrB + nnzB-1; i >= adrB; i--) { - mjtNum b = B[i]; - if (b == 0) { - continue; - } - int j = B_colind[i]; - int adrC = m->M_rowadr[j]; - mju_addToSclSparseInc(B + adrB, d->qLD + adrC, - nnzB, B_colind + adrB, - m->M_rownnz[j]-1, m->M_colind + adrC, -b); - } - - // B(r,:) <- sqrt(inv(D)) * B(r,:) - for (int i=adrB; i < adrB + nnzB; i++) { - int j = B_colind[i]; - B[i] *= sqrtInvD[j]; - } - } - - // B supernodes are identical to J supernodes - const int* B_rowsuper = d->efc_J_rowsuper; - - // construct B transposed - int* BT_rownnz = mjSTACKALLOC(d, nv, int); - int* BT_rowadr = mjSTACKALLOC(d, nv, int); - int* BT_colind = mjSTACKALLOC(d, nB, int); - mjtNum* BT = mjSTACKALLOC(d, nB, mjtNum); - mju_transposeSparse(BT, B, nefc, nv, - BT_rownnz, BT_rowadr, BT_colind, NULL, - B_rownnz, B_rowadr, B_colind); + // construct Y transposed + int* YT_rownnz = mjSTACKALLOC(d, nv, int); + int* YT_rowadr = mjSTACKALLOC(d, nv, int); + int* YT_colind = mjSTACKALLOC(d, d->nY, int); + mjtNum* YT = mjSTACKALLOC(d, d->nY, mjtNum); + mju_transposeSparse(YT, d->efc_Y, nefc, nv, + YT_rownnz, YT_rowadr, YT_colind, NULL, + d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind); // allocate AR row nonzeros and addresses on arena d->efc_AR_rownnz = mj_arenaAllocByte(d, sizeof(int) * nefc, _Alignof(int)); @@ -2891,8 +2899,8 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { int* diagind = mjSTACKALLOC(d, nefc, int); d->nA = mju_sqrMatTDSparseSymbolic( d->efc_AR_rownnz, d->efc_AR_rowadr, NULL, diagind, - nv, nefc, BT_rownnz, BT_rowadr, BT_colind, - B_rownnz, B_rowadr, B_colind, B_rowsuper, d); + nv, nefc, YT_rownnz, YT_rowadr, YT_colind, + d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind, Y_rowsuper, d); // allocate A values and column indices on arena d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum)); @@ -2905,17 +2913,18 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { return; } - // A = B * B': symbolic phase + // A = Y * Y': symbolic phase mju_sqrMatTDSparseSymbolic( d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, diagind, - nv, nefc, BT_rownnz, BT_rowadr, BT_colind, - B_rownnz, B_rowadr, B_colind, B_rowsuper, d); + nv, nefc, YT_rownnz, YT_rowadr, YT_colind, + d->efc_Y_rownnz, d->efc_Y_rowadr, d->efc_Y_colind, Y_rowsuper, d); - // A = B * B': numeric phase + // A = Y * Y': numeric phase mju_sqrMatTDSparseNumeric( d->efc_AR, nefc, d->efc_AR_rownnz, d->efc_AR_rowadr, - d->efc_AR_colind, diagind, BT, BT_rownnz, BT_rowadr, - BT_colind, B, B_rownnz, B_rowadr, B_colind, B_rowsuper, NULL, d); + d->efc_AR_colind, diagind, YT, YT_rownnz, YT_rowadr, + YT_colind, d->efc_Y, d->efc_Y_rownnz, d->efc_Y_rowadr, + d->efc_Y_colind, Y_rowsuper, NULL, d); // AR = A + diag(R) for (int i=0; i < nefc; i++) { @@ -2923,11 +2932,24 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { } } - // dense + // dense Y = backsubM2(J')' and its transpose else { - d->nA = nefc * nefc; + // arena-allocate efc_Y + d->nY = nefc * nv; + d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum)); + if (!d->efc_Y) { + mj_warning(d, mjWARN_CNSTRFULL, d->narena); + mj_clearEfc(d); + d->parena = d->ncon * sizeof(mjContact); + mj_freeStack(d); + return; + } + + // Y = backsubM2(J')' + mj_solveM2(m, d, d->efc_Y, d->efc_J, sqrtInvD, nefc); // arena-allocate efc_AR + d->nA = nefc * nefc; d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum)); if (!d->efc_AR) { mj_warning(d, mjWARN_CNSTRFULL, d->narena); @@ -2937,18 +2959,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { return; } - // space for B = backsubM2(J')' and its transpose - mjtNum* B = mjSTACKALLOC(d, nefc*nv, mjtNum); - mjtNum* BT = mjSTACKALLOC(d, nv*nefc, mjtNum); + // construct YT on stack + mjtNum* YT = mjSTACKALLOC(d, nv*nefc, mjtNum); + mju_transpose(YT, d->efc_Y, nefc, nv); - // B = backsubM2(J')' - mj_solveM2(m, d, B, d->efc_J, sqrtInvD, nefc); - - // construct BT - mju_transpose(BT, B, nefc, nv); - - // AR = B * B' - mju_sqrMatTD(d->efc_AR, BT, NULL, nv, nefc); + // AR = Y * Y' + mju_sqrMatTD(d->efc_AR, YT, NULL, nv, nefc); // add R to diagonal of AR for (int r=0; r < nefc; r++) { diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 5626baee..6a1e10fe 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -1320,6 +1320,7 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { d->nl = 0; d->nefc = 0; d->nJ = 0; + d->nY = 0; d->nA = 0; d->nisland = 0; d->nidof = 0; diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index a8f407d9..921c028a 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1192,7 +1192,7 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char* // BVHs fprintf(fp, "BVH:\n"); - fprintf(fp, " %-8s%-8s%-8s%-10s%-s\n","id", "depth", "nodeid", "child[0]" ,"child[1]"); + fprintf(fp, " %-8s%-8s%-8s%-10s%-s\n", "id", "depth", "nodeid", "child[0]", "child[1]"); for (int i=0; i < m->nbvh; i++) { fprintf(fp, " %-8d%-8d% -8d% -10d% -d\n", i, m->bvh_depth[i], m->bvh_nodeid[i], m->bvh_child[2*i], m->bvh_child[2*i+1]); @@ -1204,7 +1204,6 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char* } } - // print mjModel to text file void mj_printModel(const mjModel* m, const char* filename) { mj_printFormattedModel(m, filename, FLOAT_FORMAT); @@ -1545,7 +1544,6 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena mjtNum force[6] = {0}; mj_contactForce(m, d, i, force); printVector(" force ", force, 6, fp, float_format); - } if (d->ncon) fprintf(fp, "\n"); @@ -1568,6 +1566,11 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena d->efc_J_rowadr, d->efc_J_colind, fp, float_format); mj_printSparsity("J: constraint Jacobian", d->nefc, m->nv, d->efc_J_rowadr, NULL, d->efc_J_rownnz, d->efc_J_rowsuper, d->efc_J_colind, fp); + if (d->nY) { + mj_printSparsity("EFC_Y: inverse constraint inertia square root", d->nefc, m->nv, + d->efc_Y_rowadr, NULL, d->efc_Y_rownnz, d->efc_J_rowsuper, + d->efc_Y_colind, fp); + } if (d->nisland) { mj_printBlockSparsity("IEFC_J: block-diagonalized constraint Jacobian (nnzs are island ids)", d->nefc, d->nidof, d->nisland, @@ -1582,7 +1585,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena printArray2dInt("EFC_AR_ROWADR", d->nefc, 1, d->efc_AR_rowadr, fp); printSparse("EFC_AR", d->efc_AR, d->nefc, d->efc_AR_rownnz, d->efc_AR_rowadr, d->efc_AR_colind, fp, float_format); - mj_printSparsity("efc_AR: inverse constraint inertia", d->nefc, d->nefc, d->efc_AR_rowadr, + mj_printSparsity("EFC_AR: inverse constraint inertia", d->nefc, d->nefc, d->efc_AR_rowadr, NULL, d->efc_AR_rownnz, NULL, d->efc_AR_colind, fp); } } diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index d546c345..a85fd5bc 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5724,6 +5724,7 @@ public unsafe struct mjData_ { public int nl; public int nefc; public int nJ; + public int nY; public int nA; public int nisland; public int nidof; @@ -5883,6 +5884,10 @@ public unsafe struct mjData_ { public double* iefc_frictionloss; public double* iefc_D; public double* iefc_R; + public int* efc_Y_rownnz; + public int* efc_Y_rowadr; + public int* efc_Y_colind; + public double* efc_Y; public int* efc_AR_rownnz; public int* efc_AR_rowadr; public int* efc_AR_colind; diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index e78ece4f..1ee87963 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/wasm/codegen/generated/bindings.cc @@ -6494,6 +6494,12 @@ struct MjData { void set_nJ(int value) { ptr_->nJ = value; } + int nY() const { + return ptr_->nY; + } + void set_nY(int value) { + ptr_->nY = value; + } int nA() const { return ptr_->nA; } @@ -7005,6 +7011,18 @@ struct MjData { emscripten::val iefc_R() const { return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_R)); } + emscripten::val efc_Y_rownnz() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_Y_rownnz)); + } + emscripten::val efc_Y_rowadr() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_Y_rowadr)); + } + emscripten::val efc_Y_colind() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->nY, ptr_->efc_Y_colind)); + } + emscripten::val efc_Y() const { + return emscripten::val(emscripten::typed_memory_view(ptr_->nY, ptr_->efc_Y)); + } emscripten::val efc_AR_rownnz() const { return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->efc_AR_rownnz)); } @@ -11547,6 +11565,10 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("efc_J_rowsuper", &MjData::efc_J_rowsuper) .property("efc_KBIP", &MjData::efc_KBIP) .property("efc_R", &MjData::efc_R) + .property("efc_Y", &MjData::efc_Y) + .property("efc_Y_colind", &MjData::efc_Y_colind) + .property("efc_Y_rowadr", &MjData::efc_Y_rowadr) + .property("efc_Y_rownnz", &MjData::efc_Y_rownnz) .property("efc_aref", &MjData::efc_aref) .property("efc_b", &MjData::efc_b) .property("efc_diagApprox", &MjData::efc_diagApprox) @@ -11626,6 +11648,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) { .property("moment_rownnz", &MjData::moment_rownnz) .property("nA", &MjData::nA, &MjData::set_nA, reference()) .property("nJ", &MjData::nJ, &MjData::set_nJ, reference()) + .property("nY", &MjData::nY, &MjData::set_nY, reference()) .property("narena", &MjData::narena, &MjData::set_narena, reference()) .property("nbody_awake", &MjData::nbody_awake, &MjData::set_nbody_awake, reference()) .property("nbuffer", &MjData::nbuffer, &MjData::set_nbuffer, reference()) diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py index 111c9997..9c4832b2 100644 --- a/wasm/codegen/generators/constants.py +++ b/wasm/codegen/generators/constants.py @@ -315,6 +315,10 @@ MJDATA_SIZES: tuple[str, ...] = ( "efc_AR_colind", "efc_AR_rowadr", "efc_AR_rownnz", + "efc_Y", + "efc_Y_colind", + "efc_Y_rowadr", + "efc_Y_rownnz", "efc_D", "efc_J", "efc_JT",