Allow self-collision for interpolated flexes.

Flexes with `dof="trilinear"` or `dof="quadratic"` could not set
`selfcollide`. The vertices of interpolated flexes have no body, so the
shared-body filter that skips adjacent elements of other flexes does not apply,
and every pair of adjacent elements would collide at rest.

- Remove the compiler error for interpolated flexes with self-collision.
- Skip element pairs with vertices in the same or neighboring cells of the
  interpolation grid (`mjc_ElemSharedNode`) in all self-collision midphase
  paths (direct, SAP, BVH). Such elements share grid nodes and can only
  intersect if the grid inverts locally, so their contacts would only fight
  the elasticity. A vertex on an internal cell plane counts for the cells on
  both sides, so the rule is symmetric and robust to rounding. Self-collision
  thus requires at least 3 cells along an axis, and flexes with at most 2
  cells along every axis skip it altogether.
- BVH self- and flex-flex leaves are box-filtered before pairs are pushed, and
  marked active only when a pair is pushed.
- `cellcount` now also sizes the pinned-node array of flexcomp types grid, box,
  cylinder and ellipsoid, which could crash.
- MJWarp raises `NotImplementedError` for interpolated flexes with
  self-collision; the `ipc` mode does not handle interpolated flexes.

Self-collision is opt-in for existing models (interpolated flexes could only
compile with `selfcollide="none"`), but `auto` is the default for new ones. On
a 560-triangle box surface with `cellcount="2 2 3"`, it adds about 0.5 ms per
step at rest (SAP; 0.8 ms with BVH), and pressed flat by a plate collision takes
5.9 ms per step instead of 0.8 ms, mostly in the narrowphase.

PiperOrigin-RevId: 1003063514
Change-Id: I25d9f408639ae8c22576f2cee01020963038523a

Author: quagla

doc/XMLreference.rst

@@ -4801,6 +4801,10 @@ extensions specific to flexes.
    **sap** in 1D and 2D, and **bvh** in 3D. Which strategy performs better depends on the specifics of the model. The
    automatic setting is just a simple rule which we have found to perform well in general.
 
+   For interpolated flexes (**trilinear** or **quadratic** :ref:`dofs<body-flexcomp-dof>`), elements in the same or
+   neighboring :ref:`cells<body-flexcomp-cellcount>` do not self-collide, so self-collision requires at least 3 cells
+   along an axis.
+
 .. _flex-contact-activelayers:
 
 :at:`activelayers`: :at-val:`int(1), "1"`

doc/changelog.rst

@@ -24,6 +24,8 @@ Engine
     Models solved with PGS, noslip or islands disabled can stop at slightly different iterations. The
     :ref:`meaninertia<statistic-meaninertia>` attribute is deprecated and will be removed in a future release (compiling
     a model that sets it now gives a warning; remove the attribute to migrate).
+- Interpolated flexes (**trilinear** or **quadratic** :ref:`dofs<body-flexcomp-dof>`) now support
+  :ref:`self-collision<flex-contact-selfcollide>`, excluding elements in the same or neighboring cells.
 
 Compiler
 ^^^^^^^^

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -505,6 +505,9 @@ def _check_friction(name: str, id_: int, condim: int, friction, checks):
       raise NotImplementedError("Flex internal collisions are not implemented.")
     if (mjm.flex_rigid != 0).any():
       raise NotImplementedError("Rigid flexes are not implemented.")
+    # TODO(team): skip element pairs driven by a common grid node, as in mjc_ElemSharedNode
+    if np.any((mjm.flex_interp != 0) & (mjm.flex_selfcollide != 0) & ((mjm.flex_contype & mjm.flex_conaffinity) != 0)):
+      raise NotImplementedError("Self-collision of interpolated flexes (dof=trilinear or quadratic) is not implemented.")
   m.nmaxcondim = np.concatenate(condim_arrays).max()
   m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1))
   m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any()

src/engine/engine_collision_driver.c

@@ -544,6 +544,17 @@ static int pushElemElem(mjData* d, int f1, int e1, int f2, int e2, int group, in
 }
 
 
+// push an elem-elem self-collision pair of flex f onto the arena, unless it is excluded
+static int pushSelfElemElem(const mjModel* m, mjData* d, int f, int e1, int e2, int group,
+                            int npair) {
+  // interpolated flex: skip elements driven by a common grid node
+  if (mjc_ElemSharedNode(m, f, e1, e2)) {
+    return npair;
+  }
+  return pushElemElem(d, f, e1, f, e2, group, npair);
+}
+
+
 // filter candidate collision pair; return 0 if pair should be discarded
 static int filterCollisionPair(const mjModel* m, mjData* d, int g1, int g2, int ipair,
                                int merged, int startadr, int pairadr) {
@@ -811,8 +822,12 @@ void mj_collision(const mjModel* m, mjData* d) {
       // under the ipc flag the IPC step resolves a dim-2 flex's self-contact itself
       if (mjc_ipcOwnsFlexFlex(m, f, f)) continue;
 
-      // active element collisions
-      if (m->flex_selfcollide[f] != mjFLEXSELF_NONE) {
+      // active element collisions; for an interpolated flex with at most 2 cells along every
+      // axis, all cells are neighbors, so every pair is excluded by mjc_ElemSharedNode
+      const int* cellnum = m->flex_cellnum + 3*f;
+      int all_excluded = m->flex_interp[f] && cellnum[0] <= 2 && cellnum[1] <= 2 &&
+                         cellnum[2] <= 2;
+      if (m->flex_selfcollide[f] != mjFLEXSELF_NONE && !all_excluded) {
         int group = flexReducible(m, f) ? ngroup++ : -1;
 
         // element-element: midphase
@@ -836,7 +851,7 @@ void mj_collision(const mjModel* m, mjData* d) {
             if (mj_isElemActive(m, f, e1)) {
               for (int e2=e1+1; e2 < flex_elemnum; e2++) {
                 if (mj_isElemActive(m, f, e2)) {
-                  ncandidate = pushElemElem(d, f, e1, f, e2, group, ncandidate);
+                  ncandidate = pushSelfElemElem(m, d, f, e1, e2, group, ncandidate);
                 }
               }
             }
@@ -1130,15 +1145,26 @@ static int mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
 
     // flex : flex
     else {
-      // both are leaves
-      // box filter applied in mjc_ElemElem, bitmask filter applied earlier
+      // both are leaves: same box filter as mjc_ElemElem (margin and gap ignored in
+      // self-collisions), bitmask filter applied earlier
       if (isleaf1 && isleaf2) {
-        if (bf1 != bf2 || nodeid1 != nodeid2) {
-          npair = pushElemElem(d, f1, nodeid1, f2, nodeid2, group, npair);
+        mjtNum margin = 0;
+        if (bf1 != bf2) {
+          margin = mj_assignMargin(m, m->flex_margin[f1] + m->flex_margin[f2]) +
+                   m->flex_gap[f1] + m->flex_gap[f2];
         }
-        if (mark_active) {
-          d->bvh_active[node1 + bvhadr1] = true;
-          d->bvh_active[node2 + bvhadr2] = true;
+        if (!filterBox(d->flexelem_aabb + 6*(m->flex_elemadr[f1] + nodeid1),
+                       d->flexelem_aabb + 6*(m->flex_elemadr[f2] + nodeid2), margin)) {
+          int npair_old = npair;
+          if (bf1 != bf2) {
+            npair = pushElemElem(d, f1, nodeid1, f2, nodeid2, group, npair);
+          } else if (nodeid1 != nodeid2) {
+            npair = pushSelfElemElem(m, d, f1, nodeid1, nodeid2, group, npair);
+          }
+          if (mark_active && npair > npair_old) {
+            d->bvh_active[node1 + bvhadr1] = true;
+            d->bvh_active[node2 + bvhadr2] = true;
+          }
         }
         continue;
       }
@@ -2324,7 +2350,7 @@ static int mj_collideFlexSAPPairs(const mjModel* m, mjData* d, int f, int group,
   for (int i=0; i < nsappair; i++) {
     int e1 = elid[sappair[i].hi];
     int e2 = elid[sappair[i].lo];
-    npair = pushElemElem(d, f, e1, f, e2, group, npair);
+    npair = pushSelfElemElem(m, d, f, e1, e2, group, npair);
   }
 
   mj_freeStack(d);

src/engine/engine_collision_flex.c

@@ -386,6 +386,61 @@ int mjc_GeomElem(const mjModel* m, mjData* d, mjPreContact* con, int g, int f, i
 }
 
 
+// range [lo, hi] of the interpolation grid cells containing vertex v of flex f, per axis: a vertex
+//   on an internal cell plane belongs to the cells on both sides, so that the result does not
+//   depend on the side the vertex is rounded to; vertices within eps of a plane are treated as on
+//   the plane, where eps is relative to the cell size and covers the rounding of flex_vert0 in
+//   [0, 1] in single precision
+static void vertCellRange(const mjModel* m, int f, int v, int lo[3], int hi[3]) {
+  const mjtNum* coord = m->flex_vert0 + 3*(m->flex_vertadr[f] + v);
+  const int* cellnum = m->flex_cellnum + 3*f;
+  for (int i = 0; i < 3; i++) {
+    mjtNum u = coord[i] * cellnum[i];
+    mjtNum eps = 1e-6 * cellnum[i];
+    lo[i] = mjMAX(0, mjMIN(cellnum[i] - 1, (int)mju_floor(u - eps)));
+    hi[i] = mjMAX(0, mjMIN(cellnum[i] - 1, (int)mju_floor(u + eps)));
+  }
+}
+
+
+// return 1 if elements e1 and e2 of flex f have vertices in the same or in neighboring cells of
+//   its interpolation grid, 0 for non-interpolated flexes. Such elements are driven by a common
+//   grid node (neighboring cells share a face, edge or corner of nodes, for both trilinear and
+//   quadratic interpolation) and can only intersect if the grid inverts locally, which is resisted
+//   by elasticity; their vertices have no body for the shared-body filter
+int mjc_ElemSharedNode(const mjModel* m, int f, int e1, int e2) {
+  if (!m->flex_interp[f]) {
+    return 0;
+  }
+
+  int dim = m->flex_dim[f];
+  const int* edata1 = m->flex_elem + m->flex_elemdataadr[f] + e1*(dim + 1);
+  const int* edata2 = m->flex_elem + m->flex_elemdataadr[f] + e2*(dim + 1);
+  int lo1[4][3], hi1[4][3], lo2[4][3], hi2[4][3];
+  for (int i = 0; i <= dim; i++) {
+    vertCellRange(m, f, edata1[i], lo1[i], hi1[i]);
+    vertCellRange(m, f, edata2[i], lo2[i], hi2[i]);
+  }
+
+  // cell ranges within distance 1 along every axis
+  for (int i1 = 0; i1 <= dim; i1++) {
+    for (int i2 = 0; i2 <= dim; i2++) {
+      int shared = 1;
+      for (int k = 0; k < 3; k++) {
+        if (lo2[i2][k] - hi1[i1][k] > 1 || lo1[i1][k] - hi2[i2][k] > 1) {
+          shared = 0;
+          break;
+        }
+      }
+      if (shared) {
+        return 1;
+      }
+    }
+  }
+  return 0;
+}
+
+
 // test two elems for collision, return number of contacts
 int mjc_ElemElem(const mjModel* m, mjData* d, mjPreContact* con, int f1, int e1,
                  int f2, int e2, mjtNum margin) {

src/engine/engine_collision_flex.h

@@ -31,6 +31,10 @@ int mjc_PlaneFlex(const mjModel* m, mjData* d, mjPreContact* con, int* vert, int
 int mjc_GeomElem(const mjModel* m, mjData* d, mjPreContact* con, int g, int f, int e,
                  mjtNum margin);
 
+// return 1 if elements e1 and e2 of flex f are driven by a common node of its
+// interpolation grid
+int mjc_ElemSharedNode(const mjModel* m, int f, int e1, int e2);
+
 // test two elems for collision, return number of contacts
 int mjc_ElemElem(const mjModel* m, mjData* d, mjPreContact* con, int f1, int e1,
                  int f2, int e2, mjtNum margin);

src/user/user_flexcomp.cc

@@ -309,14 +309,12 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
   int nnode = 0;
   if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) {
     int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
-    // multi-cell count for mesh/direct/gmsh, else single cell
+    // cellcount applies to all types (see the nodal mesh below), else single cell
     int cx = 1, cy = 1, cz = 1;
-    if (type == mjFCOMPTYPE_MESH || type == mjFCOMPTYPE_DIRECT || type == mjFCOMPTYPE_GMSH) {
-      if (cellcount[0] >= 0) {
-        cx = cellcount[0];
-        cy = cellcount[1];
-        cz = cellcount[2];
-      }
+    if (cellcount[0] >= 0) {
+      cx = cellcount[0];
+      cy = cellcount[1];
+      cz = cellcount[2];
     }
     nnode = (cx * order + 1) * (cy * order + 1) * (cz * order + 1);
   }
@@ -2510,9 +2508,6 @@ void mjCFlex::Compile(const mjVFS* vfs) {
   if (edgestiffness > 0 && dim > 1) {
     throw mjCError(this, "edge stiffness only available for dim=1, please use elasticity plugins");
   }
-  if (interpolated && selfcollide != mjFLEXSELF_NONE) {
-    throw mjCError(this, "trilinear interpolation cannot do self-collision");
-  }
   nelem = (int)elem_.size() / (dim + 1);
 
   // elastic2d checks