Add sleep related data structures

PiperOrigin-RevId: 829055431
Change-Id: I1ccbd77a57044a754ae7db611b4c2c0010fbda53

src/engine/CMakeLists.txt

@@ -67,6 +67,8 @@ set(MUJOCO_ENGINE_SRCS
     engine_sensor.h
     engine_setconst.c
     engine_setconst.h
+    engine_sleep.c
+    engine_sleep.h
     engine_solver.c
     engine_solver.h
     engine_sort.h

src/engine/engine_init.c

@@ -62,6 +62,9 @@ void mj_defaultOption(mjOption* opt) {
   opt->noslip_tolerance   = 1e-6;
   opt->ccd_tolerance      = 1e-6;
 
+  // sleep settings
+  opt->sleep_tolerance    = 1e-4;
+
   // physical constants
   opt->gravity[0]         = 0;
   opt->gravity[1]         = 0;

src/engine/engine_io.c

@@ -30,6 +30,7 @@
 #include "engine/engine_macro.h"
 #include "engine/engine_memory.h"
 #include "engine/engine_plugin.h"
+#include "engine/engine_sleep.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -1361,6 +1362,16 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
     mju_copy(d->qpos, m->qpos0, m->nq);
   }
 
+  static int kAwake = -(1+mjMINAWAKE);  // tree_asleep value for fully awake tree
+
+  // set all trees to awake
+  for (int i=0; i < m->ntree; i++) {
+    d->tree_asleep[i] = kAwake;
+  }
+
+  // update sleep arrays and counters
+  mj_updateSleep(m, d);
+
   // set mocap_pos/quat = body_pos/quat for mocap bodies
   if (m->body_mocapid) {
     for (int i=0; i < m->nbody; i++) {
@@ -1606,6 +1617,8 @@ const char* mj_validateReferences(const mjModel* m) {
   X(dof_jntid,          nv,             njnt          , 0                      ) \
   X(dof_parentid,       nv,             nv            , 0                      ) \
   X(dof_Madr,           nv,             nM            , 0                      ) \
+  X(tree_bodyadr,       ntree,          nbody         , m->tree_bodynum        ) \
+  X(tree_dofadr,        ntree,          nv            , m->tree_dofnum         ) \
   X(geom_bodyid,        ngeom,          nbody         , 0                      ) \
   X(geom_matid,         ngeom,          nmat          , 0                      ) \
   X(site_bodyid,        nsite,          nbody         , 0                      ) \
@@ -1653,6 +1666,7 @@ const char* mj_validateReferences(const mjModel* m) {
   X(plugin_attradr,     nplugin,        npluginattr   , 0                      ) \
   X(tendon_adr,         ntendon,        nwrap         , m->tendon_num          ) \
   X(tendon_matid,       ntendon,        nmat          , 0                      ) \
+  X(tendon_treeid,      ntendon*2,      ntree         , 0                      ) \
   X(numeric_adr,        nnumeric,       nnumericdata  , m->numeric_size        ) \
   X(text_adr,           ntext,          ntextdata     , m->text_size           ) \
   X(tuple_adr,          ntuple,         ntupledata    , m->tuple_size          ) \

src/engine/engine_island.c

@@ -400,7 +400,7 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
 // discover islands:
 //   nisland, island_idofadr, dof_island, dof_islandnext, island_efcadr, efc_island, efc_islandnext
 void mj_island(const mjModel* m, mjData* d) {
-  int nv = m->nv, nefc = d->nefc, ntree=m->ntree;
+  int nv = m->nv, nefc = d->nefc, ntree = m->ntree, nJ = d->nJ;
 
   // no constraints: quick return
   if (mjDISABLED(mjDSBL_ISLAND) || !nefc) {
@@ -411,11 +411,11 @@ void mj_island(const mjModel* m, mjData* d) {
   mj_markStack(d);
 
   // allocate edge array, nJ is an upper bound
-  int* edge = mjSTACKALLOC(d, 2*d->nJ, int);
+  int* edge = mjSTACKALLOC(d, 2*nJ, int);
 
   // get tree-tree edges and rownnz counts from efc arrays
   int* rownnz = mjSTACKALLOC(d, ntree, int);  // number of edges per tree
-  int nedge = findEdges(m, d, rownnz, edge, d->nJ);
+  int nedge = findEdges(m, d, rownnz, edge, nJ);
 
   // compute starting address of tree's column indices while resetting rownnz
   int* rowadr = mjSTACKALLOC(d, ntree, int);
@@ -448,8 +448,10 @@ void mj_island(const mjModel* m, mjData* d) {
 
   // count nidof: total number of dofs in islands
   int nidof = 0;
-  for (int i=0; i < nv; i++) {
-    nidof += (tree_island[m->dof_treeid[i]] >= 0);
+  for (int i=0; i < ntree; i++) {
+    if (tree_island[i] >= 0) {
+      nidof += m->tree_dofnum[i];
+    }
   }
   d->nidof = nidof;
 
@@ -463,6 +465,44 @@ void mj_island(const mjModel* m, mjData* d) {
   int nisland = d->nisland;
 
 
+  // ------------------------------------- trees ---------------------------------------------------
+
+  // copy tree_island from stack to arena
+  mju_copyInt(d->tree_island, tree_island, ntree);
+
+  // compute island_ntree, number of trees per island
+  mju_zeroInt(d->island_ntree, nisland);
+  for (int i=0; i < ntree; i++) {
+    int island = tree_island[i];
+    if (island >= 0) {
+      d->island_ntree[island]++;
+    }
+  }
+
+  // compute island_itreeadr (cumsum of island_ntree)
+  d->island_itreeadr[0] = 0;
+  for (int i=1; i < nisland; i++) {
+    d->island_itreeadr[i] = d->island_itreeadr[i-1] + d->island_ntree[i-1];
+  }
+  int last_tree = d->island_itreeadr[nisland-1] + d->island_ntree[nisland-1];
+
+  // compute map_itree2tree
+  int* island_ntree2 = mjSTACKALLOC(d, nisland + 1, int);  // last elem counts unconstrained trees
+  mju_zeroInt(island_ntree2, nisland + 1);
+  for (int i=0; i < ntree; i++) {
+    int island = tree_island[i];
+    if (island >= 0) {
+      d->map_itree2tree[d->island_itreeadr[island] + island_ntree2[island]++] = i;
+    } else {
+      d->map_itree2tree[last_tree + island_ntree2[nisland]++] = i;
+    }
+  }
+
+  // SHOULD NOT OCCUR
+  if (!mju_compare(island_ntree2, d->island_ntree, nisland)) mjERROR("island_ntree miscount");
+  if (last_tree + island_ntree2[nisland] != ntree) mjERROR("miscount of unconstrained trees");
+
+
   // ------------------------------------- degrees of freedom --------------------------------------
 
   // compute dof_island, island_nv
@@ -499,7 +539,7 @@ void mj_island(const mjModel* m, mjData* d) {
     }
 
     d->map_dof2idof[dof] = idof;
-    d->map_idof2dof[idof] = dof;  // only the first ni elements of map_idof2dof are in some island
+    d->map_idof2dof[idof] = dof;  // only the first nidof elements of map_idof2dof are in some island
   }
 
   // SHOULD NOT OCCUR

src/engine/engine_print.c

@@ -665,6 +665,14 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
   }
   if (m->nv) fprintf(fp, "\n");
 
+  // trees
+  object_class = &m->ntree;
+  for (int i=0; i < m->ntree; i++) {
+    fprintf(fp, "\nTREE %d:\n", i);
+    MJMODEL_POINTERS_TREE
+  }
+  if (m->ntree) fprintf(fp, "\n");
+
   // geoms
   object_class = &m->ngeom;
   for (int i=0; i < m->ngeom; i++) {
@@ -1227,6 +1235,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
   printArray2d("ACT_DOT", m->na, 1, d->act_dot, fp, float_format);
   printArray2d("USERDATA", m->nuserdata, 1, d->userdata, fp, float_format);
   printArray2d("SENSOR", m->nsensordata, 1, d->sensordata, fp, float_format);
+  printArray2dInt("TREE_ASLEEP", m->ntree, 1, d->tree_asleep, fp);
 
   printArray2d("XPOS", m->nbody, 3, d->xpos, fp, float_format);
   printArray2d("XQUAT", m->nbody, 4, d->xquat, fp, float_format);
@@ -1312,6 +1321,13 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
     printArray2d("QHDIAGINV", m->nv, 1, d->qHDiagInv, fp, float_format);
   }
 
+  // computed sleep state
+  printArray2dInt("TREE_AWAKE", 1, m->ntree, d->tree_awake, fp);
+  printArray2dInt("BODY_AWAKE", 1, m->nbody, d->body_awake, fp);
+  printArray2dInt("BODY_AWAKE_IND", 1, d->nbody_awake, d->body_awake_ind, fp);
+  printArray2dInt("PARENT_AWAKE_IND", 1, d->nparent_awake, d->parent_awake_ind, fp);
+  printArray2dInt("DOF_AWAKE_IND", 1, d->nv_awake, d->dof_awake_ind, fp);
+
   // print qDeriv
   if (!mju_isZero(d->qDeriv, m->nD)) {
     printSparse("QDERIV", d->qDeriv, m->nv, m->D_rownnz, m->D_rowadr, m->D_colind,
@@ -1324,7 +1340,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
   }
 
   // contact
-  fprintf(fp, "CONTACT\n");
+  if (d->ncon) fprintf(fp, "CONTACT\n");
   for (int i=0; i < d->ncon; i++) {
     fprintf(fp, "  %d:\n     dim           %d\n", i, d->contact[i].dim);
     int g1 = d->contact[i].geom[0];
@@ -1454,6 +1470,11 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
   printArray2d("CFRC_EXT", m->nbody, 6, d->cfrc_ext, fp, float_format);
 
   if (d->nisland) {
+    printArray2dInt("TREE_ISLAND", 1, m->ntree, d->tree_island, fp);
+    printArray2dInt("ISLAND_NTREE", 1, d->nisland, d->island_ntree, fp);
+    printArray2dInt("ISLAND_ITREEADR", 1, d->nisland, d->island_itreeadr, fp);
+    printArray2dInt("MAP_ITREE2TREE", 1, m->ntree, d->map_itree2tree, fp);
+
     fprintf(fp, NAME_FORMAT, "DOF_ISLAND");
     for (int i = 0; i < m->nv; i++) {
       fprintf(fp, " %d", d->dof_island[i]);

src/engine/engine_setconst.c

@@ -59,6 +59,230 @@ static void mj_setM0(mjModel* m, mjData* d) {
 }
 
 
+// helper function to get the tree id of a wrap object
+static int GetWrapBodyTreeId(const mjModel* m, int wrap_index) {
+  int bodyid = -1;
+  int objid = m->wrap_objid[wrap_index];
+  switch ((mjtWrap)m->wrap_type[wrap_index]) {
+  case mjWRAP_JOINT:
+    bodyid = m->jnt_bodyid[objid];
+    break;
+  case mjWRAP_SITE:
+    bodyid = m->site_bodyid[objid];
+    break;
+  case mjWRAP_SPHERE:
+  case mjWRAP_CYLINDER:
+    bodyid = m->geom_bodyid[objid];
+    break;
+  case mjWRAP_PULLEY:
+  case mjWRAP_NONE:
+    break;
+  }
+  return (bodyid != -1) ? m->body_treeid[bodyid] : -1;
+}
+
+// set fixed quantities (do not depend on qpos0)
+static void setFixed(mjModel* m, mjData* d) {
+  mj_markStack(d);
+
+  // ----- general
+
+  // compute subtreemass
+  for (int i=0; i < m->nbody; i++) {
+    m->body_subtreemass[i] = m->body_mass[i];
+  }
+  for (int i=m->nbody-1; i > 0; i--) {
+    m->body_subtreemass[m->body_parentid[i]] += m->body_subtreemass[i];
+  }
+
+  // compute ngravcomp: number of bodies with gravity compensation
+  int ngravcomp = 0;
+  for (int i=0; i < m->nbody; i++) {
+    ngravcomp += (m->body_gravcomp[i] > 0);
+  }
+  m->ngravcomp = ngravcomp;
+
+
+  // ----- tree related (body_treeid and dof_treeid already computed)
+
+  // compute body_treeid
+  for (int i=0; i < m->nbody; i++) {
+    int weldid = m->body_weldid[i];
+    if (m->body_dofnum[weldid]) {
+      m->body_treeid[i] = m->dof_treeid[m->body_dofadr[weldid]];
+    } else {
+      m->body_treeid[i] = -1;
+    }
+  }
+
+  // compute tree_bodyadr, tree_bodynum
+  mju_zeroInt(m->tree_bodynum, m->ntree);
+  int tree_current = -1;
+  for (int i=1; i < m->nbody; i++) {
+    int treeid = m->body_treeid[i];
+    if (treeid != -1) {
+      if (treeid > tree_current) {
+        m->tree_bodyadr[++tree_current] = i;
+      }
+      m->tree_bodynum[tree_current]++;
+    }
+  }
+
+  // compute tree_dofadr, tree_dofnum
+  mju_zeroInt(m->tree_dofnum, m->ntree);
+  tree_current = -1;
+  for (int i=0; i < m->nv; i++) {
+    if (m->dof_treeid[i] > tree_current) {
+      m->tree_dofadr[++tree_current] = i;
+    }
+    m->tree_dofnum[tree_current]++;
+  }
+
+  // compute tendon_treeid, tendon_treenum
+  int* tree_marker = mjSTACKALLOC(d, m->ntree, int);  // 1 if tree has been visited, 0 otherwise
+  for (int i = 0; i < m->ntendon; i++) {
+    mju_zeroInt(tree_marker, m->ntree);
+    m->tendon_treenum[i] = 0;
+    m->tendon_treeid[2*i] = -1;
+    m->tendon_treeid[2*i+1] = -1;
+
+    for (int j = m->tendon_adr[i]; j < m->tendon_adr[i] + m->tendon_num[i]; j++) {
+      int wrap_treeid = GetWrapBodyTreeId(m, j);
+      if (wrap_treeid != -1 && !tree_marker[wrap_treeid]) {
+        tree_marker[wrap_treeid] = 1;
+        if (m->tendon_treenum[i] == 0) {
+          m->tendon_treeid[2*i] = wrap_treeid;
+        } else if (m->tendon_treenum[i] == 1) {
+          m->tendon_treeid[2*i+1] = wrap_treeid;
+        }
+        m->tendon_treenum[i]++;
+      }
+    }
+  }
+
+  // ----- apply compiler AUTO tree sleep policy
+
+  // actuators: trees with any actuated joint, site, body, or tendon do not auto-sleep
+  for (int i=0; i < m->nu; i++) {
+    int bodyid = -1;
+    int tid = m->actuator_trnid[2*i];
+    switch ((mjtTrn)m->actuator_trntype[i]) {
+    case mjTRN_JOINT:
+    case mjTRN_JOINTINPARENT:
+      bodyid = m->jnt_bodyid[tid];
+      break;
+    case mjTRN_SITE:
+    case mjTRN_SLIDERCRANK:
+      bodyid = m->site_bodyid[tid];
+      break;
+    case mjTRN_BODY:
+      bodyid = tid;
+      break;
+    case mjTRN_TENDON:
+      // wake all trees connected by this actuated tendon
+      for (int j = m->tendon_adr[tid]; j < m->tendon_adr[tid] + m->tendon_num[tid]; j++) {
+        int treeid = GetWrapBodyTreeId(m, j);
+        if (treeid != -1 && m->tree_sleep_policy[treeid] == mjSLEEP_AUTO) {
+          m->tree_sleep_policy[treeid] = mjSLEEP_AUTO_NEVER;
+        }
+      }
+      continue; // next actuator
+    case mjTRN_UNDEFINED:
+      continue; // next actuator
+    }
+
+    // wake tree containing bodyid, if any
+    if (bodyid != -1) {
+      int treeid = m->body_treeid[bodyid];
+      if (treeid != -1 && m->tree_sleep_policy[treeid] == mjSLEEP_AUTO) {
+        m->tree_sleep_policy[treeid] = mjSLEEP_AUTO_NEVER;
+      }
+    }
+  }
+
+  // trees with inter-tree tendons that have non-zero stiffness or damping do not auto-sleep
+  // if the tendon spans more than 2 trees.
+  for (int i=0; i < m->ntendon; i++) {
+    int treenum = m->tendon_treenum[i];
+
+    // tendon spans 1 or 0 trees: skip
+    if (treenum < 2) {
+      continue;
+    }
+
+    // tendon spans 2 trees and has no stiffness or damping: skip
+    if (treenum == 2 && m->tendon_stiffness[i] == 0 && m->tendon_damping[i] == 0) {
+      continue;
+    }
+
+    // tendon spans two trees with stiffness or damping or more than two trees: wake all trees
+    mju_zeroInt(tree_marker, m->ntree);
+    for (int j = m->tendon_adr[i]; j < m->tendon_adr[i] + m->tendon_num[i]; j++) {
+      int treeid = GetWrapBodyTreeId(m, j);
+
+      // if the tree is not yet marked, mark it and wake it up
+      if (treeid != -1 && !tree_marker[treeid]) {
+        tree_marker[treeid] = 1;
+        int policy = m->tree_sleep_policy[treeid];
+
+        // mark tree as never sleeping
+        if (policy == mjSLEEP_AUTO) {
+          m->tree_sleep_policy[treeid] = mjSLEEP_AUTO_NEVER;
+        }
+
+        // if the user marked it as sleepable, throw an error
+        else if (policy == mjSLEEP_ALLOWED || policy == mjSLEEP_INIT) {
+          mj_freeStack(d);
+          if (treenum > 2) {
+            mjERROR("tree %d connected to tendon %d which spans more than 2 trees, "
+                    "sleeping not allowed", treeid, i);
+          } else {
+            mjERROR("tree %d connected to tendon %d with non-zero stiffness or damping, "
+                    "sleeping not allowed", treeid, i);
+          }
+        }
+      }
+    }
+  }
+
+  // flexes: trees containing bodies that are part of any flex are not allowed to sleep
+  for (int i = 0; i < m->nflex; ++i) {
+    // node-based flex
+    if (m->flex_interp[i]) {
+      int nodenum = m->flex_nodenum[i];
+      int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[i];
+      for (int j = 0; j < nodenum; ++j) {
+        int treeid = m->body_treeid[bodyid[j]];
+        if (treeid != -1 && m->tree_sleep_policy[treeid] == mjSLEEP_AUTO) {
+          m->tree_sleep_policy[treeid] = mjSLEEP_AUTO_NEVER;
+        }
+      }
+    }
+
+    // vertex-based flex
+    else {
+      int vertnum = m->flex_vertnum[i];
+      int* bodyid = m->flex_vertbodyid + m->flex_vertadr[i];
+      for (int j = 0; j < vertnum; ++j) {
+        int treeid = m->body_treeid[bodyid[j]];
+        if (treeid != -1 && m->tree_sleep_policy[treeid] == mjSLEEP_AUTO) {
+          m->tree_sleep_policy[treeid] = mjSLEEP_AUTO_NEVER;
+        }
+      }
+    }
+  }
+
+  // set remaining trees with mjSLEEP_AUTO policy to mjSLEEP_AUTO_ALLOWED
+  for (int i = 0; i < m->ntree; i++) {
+    if (m->tree_sleep_policy[i] == mjSLEEP_AUTO) {
+      m->tree_sleep_policy[i] = mjSLEEP_AUTO_ALLOWED;
+    }
+  }
+
+  mj_freeStack(d);
+}
+
+
 // set quantities that depend on qpos0
 static void set0(mjModel* m, mjData* d) {
   int nv = m->nv;
@@ -429,6 +653,8 @@ static void setStat(mjModel* m, mjData* d) {
   mjtNum xmax[3] = {-1E+10, -1E+10, -1E+10};
   mjtNum rbound;
   mj_markStack(d);
+
+  // approximate length associated with each body
   mjtNum* body = mjSTACKALLOC(d, m->nbody, mjtNum);
 
   // compute bounding box of bodies, joint centers, geoms and sites
@@ -528,6 +754,22 @@ static void setStat(mjModel* m, mjData* d) {
     }
   }
 
+  // inherit dof length from parent body
+  for (int i=0; i < m->nv; i++) {
+    // default to linear dof, already has length units
+    m->dof_length[i] = 1;
+
+    // if rotational dof, inherit from body
+    int jnt = m->dof_jntid[i];
+    mjtJoint type = m->jnt_type[jnt];
+    int offset = i - m->jnt_dofadr[jnt];
+    if (type == mjJNT_BALL  ||
+        type == mjJNT_HINGE ||
+        (type == mjJNT_FREE && offset >= 3)) {
+      m->dof_length[i] = body[m->dof_bodyid[i]];
+    }
+  }
+
   // fix extent if too small compared to meanbody
   m->stat.extent = mju_max(m->stat.extent, 2 * m->stat.meansize);
 
@@ -553,7 +795,7 @@ static void setStat(mjModel* m, mjData* d) {
 }
 
 
-// set quantities that depend on qpos_spring
+// set quantities that depend qpos_spring
 static void setSpring(mjModel* m, mjData* d) {
   // run computations in qpos_spring
   mju_copy(d->qpos, m->qpos_spring, m->nq);
@@ -572,19 +814,18 @@ static void setSpring(mjModel* m, mjData* d) {
 }
 
 
-// entry point: set all constant fields of mjModel, except for lengthrange
+// entry point: set all remaining constant fields of mjModel, except for lengthrange
 void mj_setConst(mjModel* m, mjData* d) {
-  // compute subtreemass
-  for (int i=0; i < m->nbody; i++) {
-    m->body_subtreemass[i] = m->body_mass[i];
-  }
-  for (int i=m->nbody-1; i > 0; i--) {
-    m->body_subtreemass[m->body_parentid[i]] += m->body_subtreemass[i];
-  }
+  // set fixed quantities
+  setFixed(m, d);
 
-  // call functions
+  // set quantities that depend on qpos0
   set0(m, d);
+
+  // compute statistics
   setStat(m, d);
+
+  // set quantities that depend qpos_spring
   setSpring(m, d);
 }
 

src/engine/engine_sleep.c

@@ -0,0 +1,99 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "engine/engine_sleep.h"
+
+#include <stdio.h>
+#include <stddef.h>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+
+
+//-------------------------------- update ----------------------------------------------------------
+
+// compute sleeping arrays from tree_asleep, if flg_staticawake is set treat static bodies as awake
+void mj_updateSleepInit(const mjModel* m, mjData* d, int flg_staticawake) {
+  int ntree = m->ntree, nbody = m->nbody, nv = m->nv;
+
+  // input arrays
+  const int* tree_asleep   = d->tree_asleep;  // sleep state source of truth
+  const int* body_treeid   = m->body_treeid;
+  const int* body_parentid = m->body_parentid;
+  const int* body_mocapid  = m->body_mocapid;
+  const int* dof_bodyid    = m->dof_bodyid;
+
+  // output arrays
+  int* tree_awake        = d->tree_awake;
+  int* body_awake        = d->body_awake;
+  int* dof_awake_ind     = d->dof_awake_ind;
+  int* body_awake_ind    = d->body_awake_ind;
+  int* parent_awake_ind  = d->parent_awake_ind;
+
+  // tree_awake
+  int ntree_awake = 0;
+  for (int i=0; i < ntree; i++) {
+    tree_awake[i] = tree_asleep[i] < 0;
+    ntree_awake += tree_awake[i];
+  }
+  d->ntree_awake = ntree_awake;
+
+  // {body,parent}_awake_ind
+  int nbody_awake = 0;
+  int nparent_awake = 0;
+  for (int i=0; i < nbody; i++) {
+    // static body
+    if (body_treeid[i] < 0) {
+      if (body_mocapid[i] >= 0) {
+        // mocap body are always awake
+        body_awake[i] = mjS_AWAKE;
+      } else {
+        // mark static body unless flg_staticawake is set
+        body_awake[i] = flg_staticawake ? mjS_AWAKE : mjS_STATIC;
+      }
+    }
+
+    // dynamic body
+    else {
+      body_awake[i] = tree_awake[body_treeid[i]] ? mjS_AWAKE : mjS_ASLEEP;
+    }
+
+    // body_awake_ind: list of awake and static bodies
+    if (body_awake[i] != mjS_ASLEEP) {
+      body_awake_ind[nbody_awake++] = i;
+    }
+
+    // parent_awake_ind: list of bodies with awake or static parents
+    if (i && body_awake[body_parentid[i]] != mjS_ASLEEP) {
+      parent_awake_ind[nparent_awake++] = i;
+    }
+  }
+  d->nbody_awake = nbody_awake;
+  d->nparent_awake = nparent_awake;
+
+  // dof_awake_ind: list of awake degrees of freedom
+  int nv_awake = 0;
+  for (int i=0; i < nv; i++) {
+    int bodyid = dof_bodyid[i];
+    if (body_treeid[bodyid] >= 0 && body_awake[bodyid] == mjS_AWAKE) {
+      dof_awake_ind[nv_awake++] = i;
+    }
+  }
+  d->nv_awake = nv_awake;
+}
+
+// compute sleep arrays from tree_asleep
+void mj_updateSleep(const mjModel* m, mjData* d) {
+  mj_updateSleepInit(m, d, /*flg_staticawake*/0);
+}

src/engine/engine_sleep.h

@@ -0,0 +1,36 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_SRC_ENGINE_ENGINE_SLEEP_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_SLEEP_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// compute sleeping arrays from tree_asleep, if flg_staticawake is set, treat static bodies as awake
+MJAPI void mj_updateSleepInit(const mjModel* m, mjData* d, int flg_staticawake);
+
+// compute {ntree,nbody,nv}_awake, {tree,body}_awake, {body,dof}_awake_ind from tree_asleep
+MJAPI void mj_updateSleep(const mjModel* m, mjData* d);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_SLEEP_H_

src/engine/engine_support.c

@@ -74,7 +74,8 @@ const char* mjENABLESTRING[mjNENABLE] = {
   "Energy",
   "Fwdinv",
   "InvDiscrete",
-  "MultiCCD"
+  "MultiCCD",
+  "Sleep"
 };
 
 

src/user/user_model.cc

@@ -2931,12 +2931,25 @@ void mjCModel::CopyTree(mjModel* m) {
     }
   }
 
-  // count bodies with gravity compensation, compute ngravcomp
-  int ngravcomp = 0;
-  for (int i=0; i < nbody; i++) {
-    ngravcomp += (m->body_gravcomp[i] > 0);
+  // initialize AUTO sleep policy for all trees
+  for (int i=0; i < m->ntree; i++) {
+    m->tree_sleep_policy[i] = mjSLEEP_AUTO;
+  }
+
+  // loop over bodies, check and set non-default sleep policy
+  for (int i=1; i < nbody; i++) {
+    mjCBody* pb = bodies_[i];
+
+    // validate and set non-default sleep policy
+    if (pb->sleep != mjSLEEP_AUTO) {
+      int treeid = m->body_treeid[i];
+      // non-default sleep policy only allowed for first body in a tree
+      if (treeid == -1 || treeid == m->body_treeid[i-1]) {
+        throw mjCError(pb, "sleep policy only allowed for movable root bodies");
+      }
+      m->tree_sleep_policy[treeid] = pb->sleep;
+    }
   }
-  m->ngravcomp = ngravcomp;
 
   // recompute nM and dof_Madr given m.dof_parentid, validate
   int nM_post = 0;

src/user/user_objects.cc

@@ -2489,8 +2489,8 @@ void mjCBody::Compile(void) {
                      joints[0]->spec.type == mjJNT_FREE  &&  // it is a free joint AND
                      bodies.empty()                      &&  // no child bodies AND
                      (joints[0]->spec.align == 1 ||          // either joint.align="true"
-                      (joints[0]->spec.align == 2        &&  // or joint.align="auto"
-                       compiler->alignfree)));                  //  and compiler->align="true"
+                      (joints[0]->spec.align == 2        &&  // or (joint.align="auto"
+                       compiler->alignfree)));               //     and compiler->align="true")
 
   // free-joint alignment, phase 1 (this body + child geoms)
   double ipos_inverse[3], iquat_inverse[4];

src/xml/xml_base.h

@@ -26,6 +26,7 @@
 
 // keyword maps (defined in implementation files)
 extern const int joint_sz;
+extern const int bodysleep_sz;
 extern const int camlight_sz;
 extern const int lighttype_sz;
 extern const int integrator_sz;
@@ -50,6 +51,7 @@ extern const mjMap bool_map[];
 extern const mjMap fluid_map[];
 extern const mjMap TFAuto_map[];
 extern const mjMap joint_map[];
+extern const mjMap bodysleep_map[];
 extern const mjMap geom_map[];
 extern const mjMap camlight_map[];
 extern const mjMap lighttype_map[];

src/xml/xml_native_reader.cc

@@ -110,18 +110,18 @@ const char* MJCF[nMJCF][mjXATTRNUM] = {
             "inttotal", "interval", "tolrange"},
     {">"},
 
-    {"option", "*", "26",
+    {"option", "*", "27",
         "timestep", "impratio", "tolerance", "ls_tolerance", "noslip_tolerance",
-        "ccd_tolerance", "gravity", "wind", "magnetic", "density", "viscosity",
+        "ccd_tolerance", "sleep_tolerance", "gravity", "wind", "magnetic", "density", "viscosity",
         "o_margin", "o_solref", "o_solimp", "o_friction",
         "integrator", "cone", "jacobian",
         "solver", "iterations", "ls_iterations", "noslip_iterations", "ccd_iterations",
         "sdf_iterations", "sdf_initpoints", "actuatorgroupdisable"},
     {"<"},
-        {"flag", "?", "24", "constraint", "equality", "frictionloss", "limit", "contact",
-            "spring", "damper", "gravity", "clampctrl", "warmstart",
-            "filterparent", "actuation", "refsafe", "sensor", "midphase", "eulerdamp", "autoreset",
-            "nativeccd", "island", "override", "energy", "fwdinv", "invdiscrete", "multiccd"},
+        {"flag", "?", "25", "constraint", "equality", "frictionloss", "limit", "contact",
+            "spring", "damper", "gravity", "clampctrl", "warmstart", "filterparent", "actuation",
+            "refsafe", "sensor", "midphase", "eulerdamp", "autoreset", "nativeccd", "island",
+            "override", "energy", "fwdinv", "invdiscrete", "multiccd", "sleep"},
     {">"},
 
     {"size", "*", "14", "memory", "njmax", "nconmax", "nstack", "nuserdata", "nkey",
@@ -257,8 +257,8 @@ const char* MJCF[nMJCF][mjXATTRNUM] = {
         {"model", "*", "3", "name", "file", "content_type"},
     {">"},
 
-    {"body", "R", "11", "name", "childclass", "pos", "quat", "mocap",
-        "axisangle", "xyaxes", "zaxis", "euler", "gravcomp", "user"},
+    {"body", "R", "12", "name", "childclass", "pos", "quat", "mocap",
+        "axisangle", "xyaxes", "zaxis", "euler", "gravcomp", "sleep", "user"},
     {"<"},
         {"inertial", "?", "9", "pos", "quat", "mass", "diaginertia",
             "axisangle", "xyaxes", "zaxis", "euler", "fullinertia"},
@@ -554,6 +554,15 @@ const mjMap TFAuto_map[3] = {
 };
 
 
+// body sleep type
+const int bodysleep_sz = 4;
+const mjMap bodysleep_map[bodysleep_sz] = {
+  {"auto",          mjSLEEP_AUTO},
+  {"never",         mjSLEEP_NEVER},
+  {"allowed",       mjSLEEP_ALLOWED},
+  {"init",          mjSLEEP_INIT}
+};
+
 // joint type
 const int joint_sz = 4;
 const mjMap joint_map[joint_sz] = {
@@ -1136,6 +1145,7 @@ void mjXReader::Option(XMLElement* section, mjOption* opt) {
   ReadAttr(section, "ls_tolerance", 1, &opt->ls_tolerance, text);
   ReadAttr(section, "noslip_tolerance", 1, &opt->noslip_tolerance, text);
   ReadAttr(section, "ccd_tolerance", 1, &opt->ccd_tolerance, text);
+  ReadAttr(section, "sleep_tolerance", 1, &opt->sleep_tolerance, text);
   ReadAttr(section, "gravity", 3, opt->gravity, text);
   ReadAttr(section, "wind", 3, opt->wind, text);
   ReadAttr(section, "magnetic", 3, opt->magnetic, text);
@@ -1213,6 +1223,7 @@ void mjXReader::Option(XMLElement* section, mjOption* opt) {
     READENBL("fwdinv",      mjENBL_FWDINV)
     READENBL("invdiscrete", mjENBL_INVDISCRETE)
     READENBL("multiccd",    mjENBL_MULTICCD)
+    READENBL("sleep",       mjENBL_SLEEP)
 #undef READENBL
   }
 }
@@ -3709,8 +3720,11 @@ void mjXReader::Body(XMLElement* section, mjsBody* body, mjsFrame* frame,
       }
       ReadAlternative(elem, child->alt);
 
-      // read gravcomp
+      // gravcomp, sleep policy
       ReadAttr(elem, "gravcomp", 1, &child->gravcomp, text);
+      if (MapValue(elem, "sleep", &n, bodysleep_map, bodysleep_sz)) {
+        child->sleep = (mjtSleepPolicy) n;
+      }
 
       // read userdata
       std::vector<double> userdata;

src/xml/xml_native_writer.cc

@@ -995,6 +995,7 @@ void mjXWriter::Option(XMLElement* root) {
   WriteAttr(section, "ls_tolerance", 1, &model->option.ls_tolerance, &opt.ls_tolerance);
   WriteAttr(section, "noslip_tolerance", 1, &model->option.noslip_tolerance, &opt.noslip_tolerance);
   WriteAttr(section, "ccd_tolerance", 1, &model->option.ccd_tolerance, &opt.ccd_tolerance);
+  WriteAttr(section, "sleep_tolerance", 1, &model->option.sleep_tolerance, &opt.sleep_tolerance);
   WriteAttr(section, "gravity", 3, model->option.gravity, opt.gravity);
   WriteAttr(section, "wind", 3, model->option.wind, opt.wind);
   WriteAttr(section, "magnetic", 3, model->option.magnetic, opt.magnetic);
@@ -1067,6 +1068,7 @@ void mjXWriter::Option(XMLElement* root) {
     WRITEENBL("fwdinv",         mjENBL_FWDINV)
     WRITEENBL("invdiscrete",    mjENBL_INVDISCRETE)
     WRITEENBL("multiccd",       mjENBL_MULTICCD)
+    WRITEENBL("sleep",          mjENBL_SLEEP)
 #undef WRITEENBL
   }
 
@@ -1644,6 +1646,14 @@ void mjXWriter::Body(XMLElement* elem, mjCBody* body, mjCFrame* frame, string_vi
     if (body->gravcomp) {
       WriteAttr(elem, "gravcomp", 1, &body->gravcomp);
     }
+
+    // sleep policy
+    if (body->sleep != mjSLEEP_AUTO &&
+        body->sleep != mjSLEEP_AUTO_NEVER &&
+        body->sleep != mjSLEEP_AUTO_ALLOWED) {
+      WriteAttrKey(elem, "sleep", bodysleep_map, bodysleep_sz, body->sleep);
+    }
+
     // userdata
     WriteVector(elem, "user", body->get_userdata());