From 69c9ac074a7f54c7f9b7619b8bcc3d0eb69bfbcd Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 30 Dec 2024 15:23:41 -0800 Subject: [PATCH] Remove `mjData.qLDiagSqrtInv`, add corresponding argument to `mj_solveM2`. - `qLDiagSqrtInv` is only required for the dual solvers. It is now computed as-needed rather than unconditionally. - `mj_solveM2` now requires a new input array `sqrtInvD` which contains the square root of the inverse diagonal D (formerly saved in `qLDiagSqrtInv`). PiperOrigin-RevId: 710805133 Change-Id: I0622d6a8da3882916824e9c10bad9223c122c321 --- doc/changelog.rst | 9 +++++++++ doc/includes/references.h | 4 ++-- include/mujoco/mjdata.h | 1 - include/mujoco/mjxmacro.h | 1 - include/mujoco/mujoco.h | 3 ++- introspect/functions.py | 6 ++++++ introspect/structs.py | 8 -------- mjx/mujoco/mjx/_src/io.py | 1 - mjx/mujoco/mjx/_src/types.py | 2 -- python/mujoco/functions.cc | 8 ++++++-- python/mujoco/indexer_xmacro.h | 1 - src/engine/engine_core_constraint.c | 10 ++++++++-- src/engine/engine_core_smooth.c | 18 +++++++----------- src/engine/engine_core_smooth.h | 6 +++--- src/engine/engine_forward.c | 4 ++-- src/engine/engine_print.c | 1 - src/engine/engine_support.c | 5 ++--- src/engine/engine_vis_interact.c | 6 +++++- unity/Runtime/Bindings/MjBindings.cs | 3 +-- 19 files changed, 53 insertions(+), 44 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index c5ec24bed7..aae6bdb5b1 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -13,6 +13,15 @@ Python bindings from multiple threads. To run multiple threaded rollouts simultaneously, use the new class ``Rollout`` which encapsulates the thread pool. Contribution by :github:user:`aftersomemath`. +General +^^^^^^^ + +.. admonition:: Breaking API changes (minor) + :class: attention + + - The field ``mjData.qLDiagSqrtInv`` has been removed. This field is only required for the dual solvers. It is now + computed as-needed rather than unconditionally. Relatedly, added the corresponding argument to :ref:`mj_solveM2`. + Bug fixes ^^^^^^^^^ - Fixed a bug in the box-sphere collider, depth was incorrect for deep penetrations (:github:issue:`2206`). diff --git a/doc/includes/references.h b/doc/includes/references.h index 58a6b7d5dd..685d8e3da9 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -273,7 +273,6 @@ struct mjData_ { // computed by mj_fwdPosition/mj_factorM mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) - mjtNum* qLDiagSqrtInv; // 1/sqrt(diag(D)) (nv x 1) // computed by mj_collisionTree mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6) @@ -3221,7 +3220,8 @@ void mj_transmission(const mjModel* m, mjData* d); void mj_crb(const mjModel* m, mjData* d); void mj_factorM(const mjModel* m, mjData* d); void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); -void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); +void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, + const mjtNum* sqrtInvD, int n); void mj_comVel(const mjModel* m, mjData* d); void mj_passive(const mjModel* m, mjData* d); void mj_subtreeVel(const mjModel* m, mjData* d); diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index d718514295..c29af364c4 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -301,7 +301,6 @@ struct mjData_ { // computed by mj_fwdPosition/mj_factorM mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1) mjtNum* qLDiagInv; // 1/diag(D) (nv x 1) - mjtNum* qLDiagSqrtInv; // 1/sqrt(diag(D)) (nv x 1) // computed by mj_collisionTree mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index a54c4ff543..9edd1feda5 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -632,7 +632,6 @@ X ( mjtNum, qM, nM, 1 ) \ X ( mjtNum, qLD, nM, 1 ) \ X ( mjtNum, qLDiagInv, nv, 1 ) \ - X ( mjtNum, qLDiagSqrtInv, nv, 1 ) \ XMJV( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \ XMJV( mjtByte, bvh_active, nbvh, 1 ) \ X ( mjtNum, flexedge_velocity, nflexedge, 1 ) \ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index a6df23ad86..66f5fa9859 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -365,7 +365,8 @@ MJAPI void mj_factorM(const mjModel* m, mjData* d); MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); // Half of linear solve: x = sqrt(inv(D))*inv(L')*y -MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); +MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, + const mjtNum* sqrtInvD, int n); // Compute cvel, cdof_dot. MJAPI void mj_comVel(const mjModel* m, mjData* d); diff --git a/introspect/functions.py b/introspect/functions.py index 4e1baaadb3..b16bee9726 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -1868,6 +1868,12 @@ inner_type=ValueType(name='mjtNum', is_const=True), ), ), + FunctionParameterDecl( + name='sqrtInvD', + type=PointerType( + inner_type=ValueType(name='mjtNum', is_const=True), + ), + ), FunctionParameterDecl( name='n', type=ValueType(name='int'), diff --git a/introspect/structs.py b/introspect/structs.py index 396b4dbea2..f02158d968 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -5213,14 +5213,6 @@ doc='1/diag(D)', array_extent=('nv',), ), - StructFieldDecl( - name='qLDiagSqrtInv', - type=PointerType( - inner_type=ValueType(name='mjtNum'), - ), - doc='1/sqrt(diag(D))', - array_extent=('nv',), - ), StructFieldDecl( name='bvh_aabb_dyn', type=PointerType( diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 7ace198bea..ec00ee87ff 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -310,7 +310,6 @@ def make_data( 'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), 'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float), 'qLDiagInv': (m.nv, float) if support.is_sparse(m) else (0, float), - 'qLDiagSqrtInv': (m.nv, float), 'bvh_aabb_dyn': (m.nbvhdynamic, 6, float), 'bvh_active': (m.nbvh, jp.uint8), 'flexedge_velocity': (m.nflexedge, float), diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index a97d0cb293..8d6a62d9e4 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -1274,7 +1274,6 @@ class Data(PyTreeNode): if dense: (nv, nv) qLDiagInv: 1/diag(D) if sparse: (nv,) if dense: (0,) - qLDiagSqrtInv: 1/sqrt(diag(D)) (nv,) bvh_aabb_dyn: global bounding box (center, size) (nbvhdynamic, 6) bvh_active: volume has been added to collisions (nbvh,) flexedge_velocity: flex edge velocities (nflexedge,) @@ -1404,7 +1403,6 @@ class Data(PyTreeNode): qM: jax.Array # pylint:disable=invalid-name qLD: jax.Array # pylint:disable=invalid-name qLDiagInv: jax.Array # pylint:disable=invalid-name - qLDiagSqrtInv: jax.Array # pylint:disable=invalid-name bvh_aabb_dyn: jax.Array = _restricted_to('mujoco') bvh_active: jax.Array = _restricted_to('mujoco') # position, velocity dependent: diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index e47f5a427d..7fed04f7c9 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -234,7 +234,7 @@ PYBIND11_MODULE(_functions, pymodule) { DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM2, "n")( pymodule, [](const raw::MjModel* m, raw::MjData* d, Eigen::Ref x, - Eigen::Ref y) { + Eigen::Ref y, Eigen::Ref sqrtInvD) { if (x.rows() != y.rows()) { throw py::type_error( "the first dimension of x and y should be of the same size"); @@ -247,8 +247,12 @@ PYBIND11_MODULE(_functions, pymodule) { throw py::type_error( "the last dimension of y should be of size nv"); } + if (sqrtInvD.size() != m->nv) { + throw py::type_error( + "the size of sqrtInvD should be nv"); + } return InterceptMjErrors(::mj_solveM2)( - m, d, x.data(), y.data(), y.rows()); + m, d, x.data(), y.data(), sqrtInvD.data(), y.rows()); }); Def(pymodule); Def(pymodule); diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h index 2b2b8a430e..f39d2da2d0 100644 --- a/python/mujoco/indexer_xmacro.h +++ b/python/mujoco/indexer_xmacro.h @@ -379,7 +379,6 @@ X( mjtNum, , xaxis, njnt, 3 ) \ X( mjtNum, , cdof, nv, 6 ) \ X( mjtNum, , qLDiagInv, nv, 1 ) \ - X( mjtNum, , qLDiagSqrtInv, nv, 1 ) \ X( mjtNum, , cdof_dot, nv, 6 ) \ X( mjtNum, , qfrc_bias, nv, 1 ) \ X( mjtNum, , qfrc_passive, nv, 1 ) \ diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c index 8d380e9ed3..de206a5994 100644 --- a/src/engine/engine_core_constraint.c +++ b/src/engine/engine_core_constraint.c @@ -2067,6 +2067,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { mj_markStack(d); + // inverse square root of D from inertia LDL decomposition + mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum); + for (int i=0; i < nv; i++) { + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); + } + // space for backsubM2(J')' and its traspose mjtNum* JM2 = mjSTACKALLOC(d, nefc*nv, mjtNum); mjtNum* JM2T = mjSTACKALLOC(d, nv*nefc, mjtNum); @@ -2140,7 +2146,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { // process if not zero if (xi) { // x(i) /= sqrt(L(i,i)) - JM2[adr+i] *= d->qLDiagSqrtInv[colind[adr+i]]; + JM2[adr+i] *= sqrtInvD[colind[adr+i]]; // x(j) -= L(i,j) * x(i) int Madr_ij = m->dof_Madr[colind[adr+i]]+1; @@ -2191,7 +2197,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) { // dense else { // JM2 = backsubM2(J')' - mj_solveM2(m, d, JM2, d->efc_J, nefc); + mj_solveM2(m, d, JM2, d->efc_J, sqrtInvD, nefc); // construct JM2T mju_transpose(JM2T, JM2, nefc, nv); diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 869ea13e5e..045ec6da62 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -774,7 +774,7 @@ void mj_tendon(const mjModel* m, mjData* d) { L[i] += (mju_dist3(wpnt, wpnt+3) + wlen + mju_dist3(wpnt+6, wpnt+9))/divisor; } - // accumulate moments if consequtive points are in different bodies + // accumulate moments if consecutive points are in different bodies for (int k=0; k < (wlen < 0 ? 1 : 3); k++) { if (wbody[k] != wbody[k+1]) { // get 3D position difference, normalize @@ -1387,8 +1387,7 @@ void mj_crb(const mjModel* m, mjData* d) { // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd -void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv, - mjtNum* qLDiagSqrtInv) { +void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv) { int cnt; int Madr_kk, Madr_ki; mjtNum tmp; @@ -1445,9 +1444,6 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu for (int i=0; i < nv; i++) { mjtNum qLDi = qLD[dof_Madr[i]]; qLDiagInv[i] = 1.0/qLDi; - if (qLDiagSqrtInv) { - qLDiagSqrtInv[i] = 1.0/mju_sqrt(qLDi); - } } } @@ -1456,7 +1452,7 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd void mj_factorM(const mjModel* m, mjData* d) { TM_START; - mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv, d->qLDiagSqrtInv); + mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv); TM_ADD(mjTIMER_POS_INERTIA); } @@ -1685,10 +1681,10 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y -void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { +void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, + const mjtNum* sqrtInvD, int n) { // local copies of key variables mjtNum* qLD = d->qLD; - mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv; int* dof_Madr = m->dof_Madr; int* dof_parentid = m->dof_parentid; int nv = m->nv; @@ -1720,7 +1716,7 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) // x <- sqrt(inv(D)) * x for (int i=0; i < nv; i++) { - x[i+offset] *= qLDiagSqrtInv[i]; // x(i) /= sqrt(L(i,i)) + x[i+offset] *= sqrtInvD[i]; // x(i) /= sqrt(L(i,i)) } } } @@ -1781,7 +1777,7 @@ void mj_comVel(const mjModel* m, mjData* d) { default: // in principle we should use the new velocity to compute cdofdot, - // but it makes no difference becase crossMotion(cdof, cdof) = 0, + // but it makes no difference because crossMotion(cdof, cdof) = 0, // and using the old velocity may be more accurate numerically mju_crossMotion(cdofdot+6*j, cvel, d->cdof+6*(bda+j)); diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index a1fab753b0..f33e09ff39 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -49,8 +49,7 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d); MJAPI void mj_crb(const mjModel* m, mjData* d); // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd -MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv, - mjtNum* qLDiagSqrtInv); +MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv); // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd MJAPI void mj_factorM(const mjModel* m, mjData* d); @@ -71,7 +70,8 @@ MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, in MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island); // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y -MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); +MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, + const mjtNum* sqrtInvD, int n); //-------------------------- velocity -------------------------------------------------------------- diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index b9e1e6bfd7..08f33625f8 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -803,7 +803,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { } // factor - mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, 0); + mj_factorI(m, d, MhB, d->qH, d->qHDiagInv); } // solve @@ -986,7 +986,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) { mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM); // factorize - mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL); + mj_factorI(m, d, MhB, d->qH, d->qHDiagInv); } // solve for qacc: (qM - dt*qDeriv) * qacc = qfrc diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 4de41fca40..e627880631 100644 --- a/src/engine/engine_print.c +++ b/src/engine/engine_print.c @@ -1107,7 +1107,6 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename, } printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format); - printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format); // B sparse structure printSparsity("B: body-dof matrix", m->nbody, m->nv, d->B_rowadr, NULL, d->B_rownnz, NULL, diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 94a03de4c4..a25f874e1c 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1087,7 +1087,6 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) { int nv = m->nv; const mjtNum* qLD = d->qLD; - const mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv; const int* dofMadr = m->dof_Madr; mju_zero(res, nv); @@ -1117,9 +1116,9 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) } } - // res = sqrt(D) * res + // res *= sqrt(D) for (int i=0; i < nv; i++) { - res[i] /= qLDiagSqrtInv[i]; + res[i] *= mju_sqrt(qLD[dofMadr[i]]); } } diff --git a/src/engine/engine_vis_interact.c b/src/engine/engine_vis_interact.c index 2cb8c8a0aa..ec19a1196d 100644 --- a/src/engine/engine_vis_interact.c +++ b/src/engine/engine_vis_interact.c @@ -541,6 +541,7 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur mjtNum* jac = mjSTACKALLOC(d, 3*nv, mjtNum); mjtNum* jacM2 = mjSTACKALLOC(d, 3*nv, mjtNum); + mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum); // invalid selected body: return if (sel <= 0 || sel >= m->nbody) { @@ -554,8 +555,11 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur mju_addTo3(selpos, d->xpos+3*sel); // compute average spatial inertia at selection point + for (int i=0; i < nv; i++) { + sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]); + } mj_jac(m, d, jac, NULL, selpos, sel); - mj_solveM2(m, d, jacM2, jac, 3); + mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3); mjtNum invmass = mju_dot(jacM2+0*nv, jacM2+0*nv, nv) + mju_dot(jacM2+1*nv, jacM2+1*nv, nv) + mju_dot(jacM2+2*nv, jacM2+2*nv, nv); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 8dec655e93..92fbd9e9b8 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -4917,7 +4917,6 @@ public unsafe struct mjData_ { public double* qM; public double* qLD; public double* qLDiagInv; - public double* qLDiagSqrtInv; public double* bvh_aabb_dyn; public byte* bvh_active; public double* flexedge_velocity; @@ -6679,7 +6678,7 @@ public unsafe struct mjvSceneState_ { public static unsafe extern void mj_solveM(mjModel_* m, mjData_* d, double* x, double* y, int n); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] -public static unsafe extern void mj_solveM2(mjModel_* m, mjData_* d, double* x, double* y, int n); +public static unsafe extern void mj_solveM2(mjModel_* m, mjData_* d, double* x, double* y, double* sqrtInvD, int n); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_comVel(mjModel_* m, mjData_* d);