Allow flex sleeping

PiperOrigin-RevId: 917817500
Change-Id: Ia3bd5e52e7c2eaa3f70c81352d82c130b1d357f6

src/engine/engine_collision_driver.c

@@ -30,6 +30,7 @@
 #include "engine/engine_inline.h"
 #include "engine/engine_macro.h"
 #include "engine/engine_memory.h"
+#include "engine/engine_sleep.h"
 #include "engine/engine_sort.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
@@ -437,6 +438,7 @@ void mj_collision(const mjModel* m, mjData* d) {
 
   int nexclude = m->nexclude, npair = m->npair, nbody = m->nbody;
   int nbodyflex = m->nbody + m->nflex;
+  int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < nbody;
 
   // reset the size of the contact array and invalidate efc arrays
   d->ncon = 0;
@@ -627,6 +629,9 @@ void mj_collision(const mjModel* m, mjData* d) {
   // flex self-collisions
   for (int f=0; f < m->nflex; f++) {
     if (!m->flex_rigid[f] && (m->flex_contype[f] & m->flex_conaffinity[f])) {
+      // skip if flex is asleep
+      if (sleep_filter && mj_sleepState(m, d, mjOBJ_FLEX, f) == mjS_ASLEEP) continue;
+
       // internal collisions
       if (m->flex_internal[f]) {
         int ncon_before = d->ncon;
@@ -1403,8 +1408,9 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) {
         add_pair(m, b1, b2, &npair, bfpair, maxpair);
       }
 
-      // add all b1:flex pairs
+      // add body:flex pairs, skip if flex asleep
       for (int f=0; f < nflex; f++) {
+        if (sleep_filter && mj_sleepState(m, d, mjOBJ_FLEX, f) == mjS_ASLEEP) continue;
         add_pair(m, b1, nbody+f, &npair, bfpair, maxpair);
       }
     }
@@ -1499,6 +1505,15 @@ int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair) {
         }
       }
 
+      // flex pair: skip if neither side is dynamically awake
+      else if (sleep_filter) {
+        int awake1 = (bf1 >= nbody) ? mj_sleepState(m, d, mjOBJ_FLEX, bf1-nbody) == mjS_AWAKE
+                                    : d->body_awake[bf1] == mjS_AWAKE && m->body_treeid[bf1] >= 0;
+        int awake2 = (bf2 >= nbody) ? mj_sleepState(m, d, mjOBJ_FLEX, bf2-nbody) == mjS_AWAKE
+                                    : d->body_awake[bf2] == mjS_AWAKE && m->body_treeid[bf2] >= 0;
+        if (!awake1 && !awake2) continue;
+      }
+
       // add bodyflex pair if there is room in buffer
       add_pair(m, bf1, bf2, &npair, bfpair, maxpair);
     }

src/engine/engine_core_constraint.c

@@ -1463,7 +1463,7 @@ int mj_contactJacobian(const mjModel* m, mjData* d, const mjContact* con, int di
                   m->geom_bodyid[con->geom[side]] :
                   m->flex_vertbodyid[m->flex_vertadr[con->flex[side]] + con->vert[side]];
     }
-    // compute Jacobian differences, skipping common dofs
+    // compute Jacobian differences, skipping common DOFs
     if (dim > 3) {
       return mj_jacDifPair(m, d, chain, bid[0], bid[1], con->pos, con->pos,
                            jac1p, jac2p, jacdifp, jac1r, jac2r, jacdifr, mj_isSparse(m), 1);
@@ -2492,6 +2492,15 @@ static int mj_nc(const mjModel* m, mjData* d, int* nnz) {
           mjERROR("contact %d involves sleeping geom %d", i, asleep1 ? g1 : g2);
         }
       }
+
+      // check flex contact sides
+      for (int side = 0; side < 2; side++) {
+        if (con->geom[side] >= 0) continue;
+        int b = mj_flexBody(m, con, side);
+        if (d->body_awake[m->body_weldid[b]] == mjS_ASLEEP) {
+          mjERROR("contact %d involves sleeping flex %d", i, con->flex[side]);
+        }
+      }
     }
 
     // compute NV only if nnz requested

src/engine/engine_island.c

@@ -331,7 +331,8 @@ static int findEdges(const mjModel* m, const mjData* d,
       // unless it is a flex equality, where the tree pattern changes per dof
       if (!(efc_type == mjCNSTR_EQUALITY &&
             (m->eq_type[efc_id] == mjEQ_FLEX ||
-             m->eq_type[efc_id] == mjEQ_FLEXVERT))) {
+             m->eq_type[efc_id] == mjEQ_FLEXVERT ||
+             m->eq_type[efc_id] == mjEQ_FLEXSTRAIN))) {
         // copy tree assignment from previous constraint and continue
         efc_tree[i] = efc_tree[i-1];
         continue;
@@ -482,7 +483,7 @@ void mj_island(const mjModel* m, mjData* d) {
   // compute dof_island, island_nv
   mju_zeroInt(d->island_nv, nisland);
   for (int i=0; i < nv; i++) {
-    // assign dofs to islands
+    // assign DOFs to islands
     int island = tree_island[m->dof_treeid[i]];  // -1 if unconstrained
     d->dof_island[i] = island;
 
@@ -499,7 +500,7 @@ void mj_island(const mjModel* m, mjData* d) {
   }
 
   // compute dof <-> idof maps
-  int* island_nv2 = mjSTACKALLOC(d, nisland + 1, int);  // last element counts unconstrained dofs
+  int* island_nv2 = mjSTACKALLOC(d, nisland + 1, int);  // last element counts unconstrained DOFs
   mju_zeroInt(island_nv2, nisland + 1);
   for (int dof=0; dof < nv; dof++) {
     int island = d->dof_island[dof];

src/engine/engine_setconst.c

@@ -291,8 +291,13 @@ static void setFixed(mjModel* m, mjData* d) {
     }
   }
 
-  // flexes: trees containing bodies that are part of any flex are not allowed to sleep
+  // flexes: constraint-free trees are not allowed to sleep
   for (int i = 0; i < m->nflex; ++i) {
+    // constrained flexes are allowed to sleep
+    if (m->flex_edgeequality[i]) {
+      continue;
+    }
+
     // node-based flex
     if (m->flex_interp[i]) {
       int nodenum = m->flex_nodenum[i];

src/engine/engine_sleep.c

@@ -275,7 +275,32 @@ int mj_wake(const mjModel* m, mjData* d) {
 }
 
 
-// wake sleeping trees that touch awake trees, return number of woke trees
+// get a representative body from a flex contact side
+int mj_flexBody(const mjModel* m, const mjContact* con, int side) {
+  int f = con->flex[side];
+
+  // flex vertex contact (non-interpolated)
+  if (con->vert[side] >= 0 && m->flex_interp[f] == 0) {
+    return m->flex_vertbodyid[m->flex_vertadr[f] + con->vert[side]];
+  }
+
+  // flex element contact
+  if (con->elem[side] >= 0) {
+    if (m->flex_interp[f] == 0) {
+      int dim = m->flex_dim[f];
+      const int* edata = m->flex_elem + m->flex_elemdataadr[f] + con->elem[side]*(dim+1);
+      return m->flex_vertbodyid[m->flex_vertadr[f] + edata[0]];
+    } else {
+      return m->flex_nodebodyid[m->flex_nodeadr[f]];
+    }
+  }
+
+  // flex vertex contact (interpolated): use first node
+  return m->flex_nodebodyid[m->flex_nodeadr[f]];
+}
+
+
+// wake sleeping trees with collision contact, return number of woke trees
 int mj_wakeCollision(const mjModel* m, mjData* d) {
   int ntree = m->ntree, ncon = d->ncon, nwoke = 0;
 
@@ -287,13 +312,10 @@ int mj_wakeCollision(const mjModel* m, mjData* d) {
   for (int i=0; i < ncon; i++) {
     const mjContact* con = d->contact + i;
 
-    // only geom-geom contacts are handled
-    if (con->geom[0] < 0 || con->geom[1] < 0) {
-      continue;
-    }
+    // resolve body on each side
+    int b1 = con->geom[0] >= 0 ? m->geom_bodyid[con->geom[0]] : mj_flexBody(m, con, 0);
+    int b2 = con->geom[1] >= 0 ? m->geom_bodyid[con->geom[1]] : mj_flexBody(m, con, 1);
 
-    int b1 = m->geom_bodyid[con->geom[0]];
-    int b2 = m->geom_bodyid[con->geom[1]];
     int tree1 = m->body_treeid[b1];
     int tree2 = m->body_treeid[b2];
 
@@ -391,17 +413,59 @@ int mj_wakeEquality(const mjModel* m, mjData* d) {
         tree2 = m->body_treeid[m->site_bodyid[id2]];
       }
       break;
+
     case mjEQ_JOINT:
       tree1 = id1 >= 0 ? m->body_treeid[m->jnt_bodyid[id1]] : -1;
       tree2 = id2 >= 0 ? m->body_treeid[m->jnt_bodyid[id2]] : -1;
       break;
+
     case mjEQ_TENDON:
       mjERROR("tendon equality does not yet support sleeping");
       continue;
+
     case mjEQ_FLEX:
     case mjEQ_FLEXVERT:
-      mjERROR("flex equality does not yet support sleeping");
+    case mjEQ_FLEXSTRAIN: {
+      int f = id1;
+      int num, adr;
+      const int* bodyid;
+      if (m->flex_interp[f]) {
+        num = m->flex_nodenum[f];
+        adr = m->flex_nodeadr[f];
+        bodyid = m->flex_nodebodyid;
+      } else {
+        num = m->flex_vertnum[f];
+        adr = m->flex_vertadr[f];
+        bodyid = m->flex_vertbodyid;
+      }
+
+      // find the first awake tree, if any
+      int awake_tree = -1;
+      for (int j = 0; j < num; j++) {
+        int treeid = m->body_treeid[bodyid[adr+j]];
+        if (treeid >= 0 && d->tree_awake[treeid]) {
+          awake_tree = treeid;
+          break;
+        }
+      }
+
+      // wake sleeping island: find first sleeping tree, wakeTree wakes them all
+      if (awake_tree >= 0) {
+        int wakeval = d->tree_asleep[awake_tree];
+        for (int j = 0; j < num; j++) {
+          int treeid = m->body_treeid[bodyid[adr+j]];
+          if (treeid >= 0 && !d->tree_awake[treeid]) {
+            nwoke += mj_wakeTree(d->tree_asleep, m->ntree, treeid, wakeval);
+            #ifdef MJ_DEBUG_SLEEP
+            printf("woke tree %d due to flex equality %d at t=%g\n", treeid, i, d->time);
+            #endif
+            break;
+          }
+        }
+      }
       continue;
+    }
+
     default:
       continue;
     }
@@ -643,6 +707,7 @@ static mjtSleepState mj_equalitySleepState(const mjModel* m, const mjData* d, in
       break;
     case mjEQ_FLEX:
     case mjEQ_FLEXVERT:
+    case mjEQ_FLEXSTRAIN:
       objtype = mjOBJ_FLEX;
       break;
     default:
@@ -751,9 +816,27 @@ mjtSleepState mj_sleepState(const mjModel* m, const mjData* d, mjtObj type, int
   case mjOBJ_SENSOR:
     return mj_sensorSleepState(m, d, i);
 
-  // always awake
-  case mjOBJ_FLEX:
-    return mjS_AWAKE;
+  case mjOBJ_FLEX: {
+    // all dynamic bodies share sleep state: find and check the first one
+    int num, adr;
+    const int* bodyid;
+    if (m->flex_interp[i]) {
+      num = m->flex_nodenum[i];
+      adr = m->flex_nodeadr[i];
+      bodyid = m->flex_nodebodyid;
+    } else {
+      num = m->flex_vertnum[i];
+      adr = m->flex_vertadr[i];
+      bodyid = m->flex_vertbodyid;
+    }
+    for (int j = 0; j < num; j++) {
+      int b = bodyid[adr+j];
+      if (m->body_treeid[b] >= 0) {
+        return (mjtSleepState) d->body_awake[b];
+      }
+    }
+    return mjS_STATIC;
+  }
 
   // undefined sleep state, return AWAKE
   case mjOBJ_UNKNOWN:

src/engine/engine_sleep.h

@@ -53,6 +53,9 @@ int mj_wakeEquality(const mjModel* m, mjData* d);
 // put trees to sleep according to tolerance, return number of slept trees
 int mj_sleep(const mjModel* m, mjData* d);
 
+// get a representative body from a flex contact side
+int mj_flexBody(const mjModel* m, const mjContact* con, int side);
+
 // return sleep state of object i
 MJAPI mjtSleepState mj_sleepState(const mjModel* m, const mjData* d, mjtObj type, int i);
 

src/engine/engine_vis_visualize.c

@@ -772,6 +772,45 @@ static void addFlexGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
     thisgeom->size[0] = m->flex_radius[i];
     setMaterial(m, thisgeom, m->flex_matid[i], m->flex_rgba+4*i, vopt->flags);
 
+    // override if visualizing islands
+    if (vopt->flags[mjVIS_ISLAND]) {
+      // find first dynamic body in flex
+      int bodyid = -1;
+      if (m->flex_interp[i]) {
+        int nodeadr = m->flex_nodeadr[i];
+        for (int j=0; j < m->flex_nodenum[i] && bodyid < 0; j++) {
+          int b = m->flex_nodebodyid[nodeadr+j];
+          if (m->body_treeid[b] >= 0) bodyid = b;
+        }
+      } else {
+        int vertadr = m->flex_vertadr[i];
+        for (int j=0; j < m->flex_vertnum[i] && bodyid < 0; j++) {
+          int b = m->flex_vertbodyid[vertadr+j];
+          if (m->body_treeid[b] >= 0) bodyid = b;
+        }
+      }
+
+      if (bodyid >= 0) {
+        // strip material
+        thisgeom->matid = -1;
+
+        int weld_id = m->body_weldid[bodyid];
+        int dof = m->body_dofadr[weld_id];
+        int island = d->nisland ? d->dof_island[dof] : -1;
+        int h = island >= 0 ? d->island_dofadr[island] : -1;
+        int awake = d->body_awake[bodyid];
+
+        // if sleep is enabled, color by first tree dof
+        if (h == -1 && mjENABLED(mjENBL_SLEEP)) {
+          int tree = m->dof_treeid[dof];
+          if (!awake) tree = mj_sleepCycle(d->tree_asleep, m->ntree, tree);
+          h = m->tree_dofadr[tree];
+        }
+
+        islandColor(thisgeom->rgba, h, awake);
+      }
+    }
+
     // set texcoord
     if (m->flex_texcoordadr[i] >= 0) {
       thisgeom->texcoord = 1;