From c0096b51a477aeff029a87afc0cda41e60e3651f Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sat, 16 Aug 2025 09:15:52 -0700 Subject: [PATCH] Move `ntree` to `mjModel` allocator sizes PiperOrigin-RevId: 795855203 Change-Id: Idab185294d53dced9ebf073b782969dc5643b95a --- doc/includes/references.h | 2 +- include/mujoco/mjmodel.h | 2 +- include/mujoco/mjxmacro.h | 2 +- python/mujoco/introspect/structs.py | 10 ++++---- src/engine/engine_io.c | 9 +++---- src/engine/engine_io.h | 4 ++-- src/user/user_model.cc | 36 ++++++++++++++++++++++++---- src/user/user_model.h | 1 + unity/Runtime/Bindings/MjBindings.cs | 2 +- 9 files changed, 48 insertions(+), 20 deletions(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 4fb0c89d..daed1b23 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -964,6 +964,7 @@ struct mjModel_ { int nbvhdynamic; // number of dynamic bounding volumes (aabb stored in mjData) int noct; // number of total octree cells in all meshes int njnt; // number of joints + int ntree; // number of kinematic trees under world body int nM; // number of non-zeros in sparse inertia matrix int nB; // number of non-zeros in sparse body-dof matrix int nC; // number of non-zeros in sparse reduced dof-dof matrix @@ -1032,7 +1033,6 @@ struct mjModel_ { // sizes set after mjModel construction int nnames_map; // number of slots in the names hash map int nJmom; // number of non-zeros in sparse actuator_moment 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 (legacy) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 4c355513..436a01f6 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -654,6 +654,7 @@ struct mjModel_ { int nbvhdynamic; // number of dynamic bounding volumes (aabb stored in mjData) int noct; // number of total octree cells in all meshes int njnt; // number of joints + int ntree; // number of kinematic trees under world body int nM; // number of non-zeros in sparse inertia matrix int nB; // number of non-zeros in sparse body-dof matrix int nC; // number of non-zeros in sparse reduced dof-dof matrix @@ -722,7 +723,6 @@ struct mjModel_ { // sizes set after mjModel construction int nnames_map; // number of slots in the names hash map int nJmom; // number of non-zeros in sparse actuator_moment 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 (legacy) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 08aca3c1..13d5b04c 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -77,6 +77,7 @@ X( nbvhdynamic ) \ X( noct ) \ X( njnt ) \ + X( ntree ) \ X( nM ) \ X( nB ) \ X( nC ) \ @@ -143,7 +144,6 @@ X( npaths ) \ X( nnames_map ) \ X( nJmom ) \ - X( ntree ) \ X( ngravcomp ) \ X( nemax ) \ X( njmax ) \ diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index 71de6e3e..e658fedf 100644 --- a/python/mujoco/introspect/structs.py +++ b/python/mujoco/introspect/structs.py @@ -883,6 +883,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of joints', ), + StructFieldDecl( + name='ntree', + type=ValueType(name='int'), + doc='number of kinematic trees under world body', + ), StructFieldDecl( name='nM', type=ValueType(name='int'), @@ -1213,11 +1218,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of non-zeros in sparse actuator_moment matrix', ), - StructFieldDecl( - name='ntree', - type=ValueType(name='int'), - doc='number of kinematic trees under world body', - ), StructFieldDecl( name='ngravcomp', type=ValueType(name='int'), diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 3adb83bc..bfda1e4d 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -463,7 +463,7 @@ static void freeModelBuffers(mjModel* m) { // allocate and initialize mjModel structure void mj_makeModel(mjModel** dest, int nq, int nv, int nu, int na, int nbody, int nbvh, - int nbvhstatic, int nbvhdynamic, int noct, int njnt, + int nbvhstatic, int nbvhdynamic, int noct, int njnt, int ntree, int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, @@ -505,6 +505,7 @@ void mj_makeModel(mjModel** dest, m->nbvhdynamic = nbvhdynamic; m->noct = noct; m->njnt = njnt; + m->ntree = ntree; m->nM = nM; m->nB = nB; m->nC = nC; @@ -648,8 +649,8 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) { if (!dest) { mj_makeModel( &dest, src->nq, src->nv, src->nu, src->na, src->nbody, src->nbvh, - src->nbvhstatic, src->nbvhdynamic, src->noct, src->njnt, src->nM, - src->nB, src->nC, src->nD, src->ngeom, src->nsite, src->ncam, + src->nbvhstatic, src->nbvhdynamic, src->noct, src->njnt, src->ntree, + src->nM, src->nB, src->nC, src->nD, src->ngeom, src->nsite, src->ncam, src->nlight, src->nflex, src->nflexnode, src->nflexvert, src->nflexedge, src->nflexelem, src->nflexelemdata, src->nflexelemedge, src->nflexshelldata, src->nflexevpair, src->nflextexcoord, src->nmesh, @@ -845,7 +846,7 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) { ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55], ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62], ints[63], ints[64], ints[65], ints[66], ints[67], ints[68], ints[69], - ints[70], ints[71], ints[72], ints[73]); + ints[70], ints[71], ints[72], ints[73], ints[74]); // read mjModel mjtSize fields mjtSize sizes[8]; diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index 2b4be47f..5faa8b68 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -53,8 +53,8 @@ void mj_defaultStatistic(mjStatistic* stat); // allocate mjModel void mj_makeModel(mjModel** dest, int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic, int noct, - int njnt, int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight, int nflex, - int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata, + int njnt, int ntree, int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight, + int nflex, int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface, int nmeshgraph, int nmeshpoly, int nmeshpolyvert, int nmeshpolymap, int nskin, int nskinvert, int nskintexvert, int nskinface, diff --git a/src/user/user_model.cc b/src/user/user_model.cc index aef2063d..c9d4c05b 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -2088,12 +2088,33 @@ void mjCModel::SetSizes() { ntuple = (int)tuples_.size(); nkey = (int)keys_.size(); nplugin = (int)plugins_.size(); - nq = nv = nu = na = nmocap = 0; + nq = nv = ntree = nu = na = nmocap = 0; - // nq, nv + // nq, nv, ntree for (int i=0; i < njnt; i++) { nq += joints_[i]->nq(); nv += joints_[i]->nv(); + + // increment ntree if this is the first joint in a moving body with static ancestry + mjCBody* parent = joints_[i]->GetParent(); + bool is_first_joint = joints_[i] == parent->joints[0]; + if (is_first_joint) { + // check if all ancestors are static + bool static_ancestry = true; + mjCBody* ancestor = parent; + while (ancestor != bodies_[0]) { + ancestor = ancestor->GetParent(); + if (!ancestor->joints.empty()) { + static_ancestry = false; + break; + } + } + + // if all ancestors are static, this joint starts a new kinematic tree + if (static_ancestry) { + ntree++; + } + } } // nu, na @@ -2899,7 +2920,12 @@ void mjCModel::CopyTree(mjModel* m) { } m->dof_treeid[i] = ntree - 1; } - m->ntree = ntree; + + // check number of trees constructed, SHOULD NOT OCCUR + if (ntree != m->ntree) { + throw mjCError(0, "unexpected number of TREEs. Counted %d, expected %d", + nullptr, ntree, m->ntree); + } // compute body_treeid for (int i=0; i < nbody; i++) { @@ -4770,8 +4796,8 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // create low-level model mj_makeModel(&m, - nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, noct, njnt, nM, nB, nC, nD, - ngeom, nsite, ncam, nlight, nflex, nflexnode, nflexvert, nflexedge, nflexelem, + nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, noct, njnt, ntree, nM, nB, nC, + nD, ngeom, nsite, ncam, nlight, nflex, nflexnode, nflexvert, nflexedge, nflexelem, nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord, nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph, nmeshpoly, nmeshpolyvert, nmeshpolymap, nskin, nskinvert, nskintexvert, nskinface, nskinbone, diff --git a/src/user/user_model.h b/src/user/user_model.h index bdbf93b3..e6a5ae9f 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -84,6 +84,7 @@ class mjCModel_ : public mjsElement { int nv; // number of degrees of freedom = dim(qvel) int nu; // number of actuators/controls int na; // number of activation variables + int ntree; // number of trees int nbvh; // number of total boundary volume hierarchies int nbvhstatic; // number of static boundary volume hierarchies int nbvhdynamic; // number of dynamic boundary volume hierarchies diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 78a3b50e..cdc0c118 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5254,6 +5254,7 @@ public unsafe struct mjModel_ { public int nbvhdynamic; public int noct; public int njnt; + public int ntree; public int nM; public int nB; public int nC; @@ -5320,7 +5321,6 @@ public unsafe struct mjModel_ { public int npaths; public int nnames_map; public int nJmom; - public int ntree; public int ngravcomp; public int nemax; public int njmax;