From 77d040073377bd586ea946446ecba73d0cc2fd2e Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 6 Oct 2023 12:50:37 -0700 Subject: [PATCH] Accelerate inertia computation for simple bodies. PiperOrigin-RevId: 571411319 Change-Id: I82732e95795cc0f0645f8273a1b084f2a9187a75 --- doc/includes/references.h | 2 +- include/mujoco/mjmodel.h | 2 +- introspect/structs.py | 2 +- src/engine/engine_setconst.c | 122 +++++++++++++++++++----------- src/user/user_model.cc | 11 +++ test/engine/engine_solver_test.cc | 2 +- 6 files changed, 92 insertions(+), 49 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 3a447e74..8133a71e 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -909,7 +909,7 @@ struct mjModel_ { int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1) int* body_geomnum; // number of geoms (nbody x 1) int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1) - mjtByte* body_simple; // body is simple (has diagonal M) (nbody x 1) + mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1) mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1) mjtNum* body_pos; // position offset rel. to parent body (nbody x 3) mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index cc06c92b..e0961f2c 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -638,7 +638,7 @@ struct mjModel_ { int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1) int* body_geomnum; // number of geoms (nbody x 1) int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1) - mjtByte* body_simple; // body is simple (has diagonal M) (nbody x 1) + mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1) mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1) mjtNum* body_pos; // position offset rel. to parent body (nbody x 3) mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4) diff --git a/introspect/structs.py b/introspect/structs.py index 04a8a5e8..e85ee1f3 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -1253,7 +1253,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtByte'), ), - doc='body is simple (has diagonal M) (nbody x 1)', + doc='1: diagonal M; 2: diag M, no rotations (nbody x 1)', ), StructFieldDecl( name='body_sameframe', diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c index 397ce15f..542c9c87 100644 --- a/src/engine/engine_setconst.c +++ b/src/engine/engine_setconst.c @@ -59,7 +59,7 @@ static void mj_setM0(mjModel* m, mjData* d) { // set quantities that depend on qpos0 static void set0(mjModel* m, mjData* d) { - int id, id1, id2, dnum, nv = m->nv; + int nv = m->nv; mjtNum A[36] = {0}, pos[3], quat[4]; mj_markStack(d); mjtNum* jac = mj_stackAllocNum(d, 6*nv); @@ -112,53 +112,85 @@ static void set0(mjModel* m, mjData* d) { // compute body_invweight0 m->body_invweight0[0] = m->body_invweight0[1] = 0.0; - for (int i=1; i < m->nbody; i++) { - if (nv) { - // inverse spatial inertia: A = J*inv(M)*J' - mj_jacBodyCom(m, d, jac, jac+3*nv, i); - mj_solveM(m, d, tmp, jac, 6); - mju_mulMatMatT(A, jac, tmp, 6, nv, 6); + for (int i=1; inbody; i++) { + // static bodies: zero invweight0 + if (m->body_weldid[i] == 0) { + m->body_invweight0[2*i] = m->body_invweight0[2*i+1] = 0; } - // average diagonal and assign - m->body_invweight0[2*i] = (A[0] + A[7] + A[14])/3; - m->body_invweight0[2*i+1] = (A[21] + A[28] + A[35])/3; + // accelerate simple bodies with no rotations + else if (m->body_simple[i]==2) { + mjtNum mass = m->body_mass[i]; + if (!mass) { // SHOULD NOT OCCUR + mjERROR("moving body %d has 0 mass", i); + } + m->body_invweight0[2*i+0] = 1/mju_max(mjMINVAL, mass); + m->body_invweight0[2*i+1] = 0; + } + + // general body: full inertia + else { + if (nv) { + // inverse spatial inertia: A = J*inv(M)*J' + mj_jacBodyCom(m, d, jac, jac+3*nv, i); + mj_solveM(m, d, tmp, jac, 6); + mju_mulMatMatT(A, jac, tmp, 6, nv, 6); + } + + // average diagonal and assign + m->body_invweight0[2*i] = (A[0] + A[7] + A[14])/3; + m->body_invweight0[2*i+1] = (A[21] + A[28] + A[35])/3; + } } // compute dof_invweight0 - for (int i=0; i < m->njnt; i++) { - id = m->jnt_dofadr[i]; - - // get number of components - if (m->jnt_type[i] == mjJNT_FREE) { - dnum = 6; - } else if (m->jnt_type[i] == mjJNT_BALL) { - dnum = 3; - } else { - dnum = 1; - } - - // inverse joint inertia: A = J*inv(M)*J' - if (nv) { - mju_zero(jac, dnum*nv); - for (int j=0; j < dnum; j++) { - jac[j*(nv+1) + id] = 1; + for (int i=0; injnt; i++) { + // simple body with no rotations: no off-diagonal inertia + if (m->body_simple[m->jnt_bodyid[i]] == 2) { + int id = m->jnt_dofadr[i]; + int bi = m->jnt_bodyid[i]; + mjtNum mass = m->body_mass[bi]; + if (!mass) { // SHOULD NOT OCCUR + mjERROR("moving body %d has 0 mass", bi); } - mj_solveM(m, d, tmp, jac, dnum); - mju_mulMatMatT(A, jac, tmp, dnum, nv, dnum); + m->dof_invweight0[id] = 1/mju_max(mjMINVAL, mass); } - // average diagonal and assign - if (dnum == 6) { - m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] = - (A[0] + A[7] + A[14])/3; - m->dof_invweight0[id+3] = m->dof_invweight0[id+4] = m->dof_invweight0[id+5] = - (A[21] + A[28] + A[35])/3; - } else if (dnum == 3) - m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] = - (A[0] + A[4] + A[8])/3; + // general joint: full inertia else { - m->dof_invweight0[id] = A[0]; + int dnum, id = m->jnt_dofadr[i]; + + // get number of components + if (m->jnt_type[i] == mjJNT_FREE) { + dnum = 6; + } else if (m->jnt_type[i] == mjJNT_BALL) { + dnum = 3; + } else { + dnum = 1; + } + + // inverse joint inertia: A = J*inv(M)*J' + if (nv) { + mju_zero(jac, dnum*nv); + for (int j=0; j < dnum; j++) { + jac[j*(nv+1) + id] = 1; + } + mj_solveM(m, d, tmp, jac, dnum); + mju_mulMatMatT(A, jac, tmp, dnum, nv, dnum); + } + + // average diagonal and assign + if (dnum == 6) { + m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] = + (A[0] + A[7] + A[14])/3; + m->dof_invweight0[id+3] = m->dof_invweight0[id+4] = m->dof_invweight0[id+5] = + (A[21] + A[28] + A[35])/3; + } else if (dnum == 3) { + m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] = + (A[0] + A[4] + A[8])/3; + } else { + m->dof_invweight0[id] = A[0]; + } } } @@ -198,8 +230,8 @@ static void set0(mjModel* m, mjData* d) { // compute missing eq_data for body constraints for (int i=0; i < m->neq; i++) { // get ids - id1 = m->eq_obj1id[i]; - id2 = m->eq_obj2id[i]; + int id1 = m->eq_obj1id[i]; + int id2 = m->eq_obj2id[i]; // connect constraint if (m->eq_type[i] == mjEQ_CONNECT) { @@ -239,8 +271,8 @@ static void set0(mjModel* m, mjData* d) { // camera compos0, pos0, mat0 for (int i=0; i < m->ncam; i++) { // get body ids - id = m->cam_bodyid[i]; // camera body - id1 = m->cam_targetbodyid[i]; // target body + int id = m->cam_bodyid[i]; // camera body + int id1 = m->cam_targetbodyid[i]; // target body // compute positional offsets mju_sub3(m->cam_pos0+3*i, d->cam_xpos+3*i, d->xpos+3*id); @@ -253,8 +285,8 @@ static void set0(mjModel* m, mjData* d) { // light compos0, pos0, dir0 for (int i=0; i < m->nlight; i++) { // get body ids - id = m->light_bodyid[i]; // light body - id1 = m->light_targetbodyid[i]; // target body + int id = m->light_bodyid[i]; // light body + int id1 = m->light_targetbodyid[i]; // target body // compute positional offsets mju_sub3(m->light_pos0+3*i, d->light_xpos+3*i, d->xpos+3*id); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index e10cf370..47f0bb53 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -1520,6 +1520,17 @@ void mjCModel::CopyTree(mjModel* m) { qposadr += nPOS[pj->type]; } + // simple body with no rotational dofs: promote to simple level 2 + if (m->body_simple[i]) { + m->body_simple[i] = 2; + for (int j=0; j<(int)pb->joints.size(); j++) { + if (pb->joints[j]->type!=mjJNT_SLIDE) { + m->body_simple[i] = 1; + break; + } + } + } + // loop over geoms for this body for (int j=0; j<(int)pb->geoms.size(); j++) { // get pointer and id diff --git a/test/engine/engine_solver_test.cc b/test/engine/engine_solver_test.cc index 330542b6..47a642ed 100644 --- a/test/engine/engine_solver_test.cc +++ b/test/engine/engine_solver_test.cc @@ -107,7 +107,7 @@ TEST_F(SolverTest, OneBigIsland) { mjData* data_noisland = mj_makeData(model); int nv = model->nv; - mjtNum tol = 1e-9; + mjtNum tol = 1e-8; // save current (default) iterations int iterations_default = model->opt.iterations;