Move ntree to mjModel allocator sizes

PiperOrigin-RevId: 795855203
Change-Id: Idab185294d53dced9ebf073b782969dc5643b95a

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];

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,

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,

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