Move ntree to mjModel allocator sizes

PiperOrigin-RevId: 795855203
Change-Id: Idab185294d53dced9ebf073b782969dc5643b95a

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)

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)

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 )              \

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'),

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->nbvhstatic, src->nbvhdynamic, src->noct, src->njnt, src->ntree,
-        src->nB, src->nC, src->nD, src->ngeom, src->nsite, src->ncam,
+        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];

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 njnt, int ntree, int nM, int nB, int nC, int nD, int ngeom, int nsite, int ncam, int nlight,
-    int nflexnode, int nflexvert, int nflexedge, int nflexelem, int nflexelemdata,
+    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,

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,
+               nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, noct, njnt, ntree, nM, nB, nC,
-               ngeom, nsite, ncam, nlight, nflex, nflexnode, nflexvert, nflexedge, nflexelem,
+               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,

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

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;