From 8ebc9ed2c8f3602aa19e91b3ea4f5ec5433ae162 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 22 Apr 2024 10:44:33 -0700 Subject: [PATCH] Add `mjModel.ngravcomp`, number of bodies with nonzero gravcomp attribute. PiperOrigin-RevId: 627087335 Change-Id: Ie85ae4150962463b8a8c8435fec068c44877aa16 --- doc/includes/references.h | 1 + include/mujoco/mjmodel.h | 1 + include/mujoco/mjxmacro.h | 1 + introspect/structs.py | 5 ++++ src/engine/engine_forward.c | 41 ++++++++++++++-------------- src/engine/engine_passive.c | 2 +- src/user/user_model.cc | 7 +++++ unity/Runtime/Bindings/MjBindings.cs | 1 + 8 files changed, 37 insertions(+), 22 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index ebf3252d..1be65826 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -922,6 +922,7 @@ struct mjModel_ { int nD; // number of non-zeros in sparse dof-dof matrix int nB; // number of non-zeros in sparse body-dof matrix int ntree; // number of kinematic trees under world body + int ngravcomp; // number of bodies with nonzero gravcomp int nemax; // number of potential equality-constraint rows int njmax; // number of available rows in constraint Jacobian int nconmax; // number of potential contacts in contact list diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 79ec6ad9..9bad4ef6 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -637,6 +637,7 @@ struct mjModel_ { int nD; // number of non-zeros in sparse dof-dof matrix int nB; // number of non-zeros in sparse body-dof matrix int ntree; // number of kinematic trees under world body + int ngravcomp; // number of bodies with nonzero gravcomp int nemax; // number of potential equality-constraint rows int njmax; // number of available rows in constraint Jacobian int nconmax; // number of potential contacts in contact list diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index daf1d30b..9cea2bc6 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -139,6 +139,7 @@ X ( njmax ) \ X ( nconmax ) \ XMJV( ntree ) \ + X ( ngravcomp ) \ X ( nuserdata ) \ XMJV( nsensordata ) \ X ( npluginstate ) \ diff --git a/introspect/structs.py b/introspect/structs.py index 45199a98..3d7e9cf5 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -1200,6 +1200,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of kinematic trees under world body', ), + StructFieldDecl( + name='ngravcomp', + type=ValueType(name='int'), + doc='number of bodies with nonzero gravcomp', + ), StructFieldDecl( name='nemax', type=ValueType(name='int'), diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 8275f751..fb017fba 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -256,6 +256,23 @@ static void clampVec(mjtNum* vec, const mjtNum* range, const mjtByte* limited, i +// return number of dofs given joint type +static int jnt_dofnum(mjtJoint type) { + switch (type) { + case mjJNT_HINGE: + case mjJNT_SLIDE: + return 1; + + case mjJNT_BALL: + return 3; + + case mjJNT_FREE: + return 6; + } +} + + + // (qpos, qvel, ctrl, act) => (qfrc_actuator, actuator_force, act_dot) void mj_fwdActuation(const mjModel* m, mjData* d) { TM_START; @@ -469,7 +486,7 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { mju_mulMatTVec(d->qfrc_actuator, moment, force, nu, nv); // actuator-level gravity compensation - if (!mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) { + if (m->ngravcomp && !mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) { int njnt = m->njnt; for (int i=0; i < njnt; i++) { // skip if gravcomp added as passive force @@ -477,28 +494,10 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { continue; } - // get number of dofs for this joint - int dofnum; - switch (m->jnt_type[i]) { - case mjJNT_HINGE: - case mjJNT_SLIDE: - dofnum = 1; - break; - - case mjJNT_BALL: - dofnum = 3; - break; - - case mjJNT_FREE: - dofnum = 6; - break; - } - // add gravcomp force + int dofnum = jnt_dofnum(m->jnt_type[i]); int dofadr = m->jnt_dofadr[i]; - for (int j=0; j < dofnum; j++) { - d->qfrc_actuator[dofadr+j] += d->qfrc_gravcomp[dofadr+j]; - } + mju_addTo(d->qfrc_actuator + dofadr, d->qfrc_gravcomp + dofadr, dofnum); } } diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index 6df4ef23..6fae24c6 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -175,7 +175,7 @@ static void mj_springdamper(const mjModel* m, mjData* d) { // body-level gravity compensation, return 1 if any, 0 otherwise static int mj_gravcomp(const mjModel* m, mjData* d) { - if (mjDISABLED(mjDSBL_GRAVITY) || mju_norm3(m->opt.gravity) == 0) { + if (!m->ngravcomp || mjDISABLED(mjDSBL_GRAVITY) || mju_norm3(m->opt.gravity) == 0) { return 0; } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 5abf4163..1497e7d8 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -2105,6 +2105,13 @@ void mjCModel::CopyTree(mjModel* m) { } } + // count bodies with gravity compensation, compute ngravcomp + int ngravcomp = 0; + for (int i=0; ibody_gravcomp[i] > 0); + } + m->ngravcomp = ngravcomp; + // compute nM and dof_Madr nM = 0; for (int i=0; i