Add vertex equality option.

PiperOrigin-RevId: 860299454
Change-Id: I7befd8e5f4c8bac6d74684a96a9cc3cee9ae1e29

src/engine/engine_core_constraint.c

@@ -376,6 +376,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
   int issparse = mj_isSparse(m), nv = m->nv;
   int id[2], size, NV, NV2, *chain = NULL, *chain2 = NULL, *buf_ind = NULL;
   int flex_edgeadr, flex_edgenum;
+  int flex_vertadr, flex_vertnum;
   mjtNum cpos[6], pos[2][3], ref[2], dif, deriv;
   mjtNum quat[4], quat1[4], quat2[4], quat3[4], axis[3];
   mjtNum *jac[2], *jacdif, *data, *sparse_buf = NULL;
@@ -643,6 +644,32 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
       }
       break;
 
+    case mjEQ_FLEXVERT:
+      // add two constraints per vertex
+      flex_vertadr = m->flex_vertadr[id[0]];
+      flex_vertnum = m->flex_vertnum[id[0]];
+      for (int v=flex_vertadr; v < flex_vertadr+flex_vertnum; v++) {
+        for (int j=0; j < 2; j++) {
+          cpos[0] = d->flexvert_length[2*v+j];
+          int row = 2*v+j;
+          if (issparse) {
+            mj_addConstraint(m, d, d->flexvert_J + m->flexvert_J_rowadr[row],
+                             cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
+                             m->flexvert_J_rownnz[row],
+                             m->flexvert_J_colind + m->flexvert_J_rowadr[row]);
+          } else {
+            mju_zero(jac[0], nv);  // reuse first row of jac[0]
+            int rowadr = m->flexvert_J_rowadr[row];
+            int rownnz = m->flexvert_J_rownnz[row];
+            for (int k=0; k<rownnz; k++) {
+              jac[0][m->flexvert_J_colind[rowadr+k]] = d->flexvert_J[rowadr+k];
+            }
+            mj_addConstraint(m, d, jac[0], cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
+          }
+        }
+      }
+      break;
+
     default:                    // SHOULD NOT OCCUR
       mjERROR("invalid equality constraint type %d", m->eq_type[i]);
     }
@@ -1138,8 +1165,23 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
         i--;
         break;
 
+      case mjEQ_FLEXVERT:
+        // process all vertices for this flex
+        f = m->eq_obj1id[id];
+        int vertadr = m->flex_vertadr[f];
+        int vertnum = m->flex_vertnum[f];
+        for (int v=vertadr; v<vertadr+vertnum; v++) {
+          int bodyid = m->flex_vertbodyid[v];
+          dA[i++] = m->body_invweight0[2*bodyid];
+          dA[i++] = m->body_invweight0[2*bodyid];
+        }
+
+        // adjust constraint counter
+        i--;
+        break;
+
       default:
-        mjERROR("unknown constraint type type %d", d->efc_type[i]);    // SHOULD NOT OCCUR
+        mjERROR("unknown constraint type %d", d->efc_type[i]);    // SHOULD NOT OCCUR
       }
       break;
 
@@ -1605,7 +1647,7 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
   int nv = m->nv, neq = m->neq;
   int id[2], size, NV, NV2, *chain = NULL, *chain2 = NULL;
   int issparse = (nnz != NULL);
-  int flex_edgeadr, flex_edgenum;
+  int flex_edgeadr, flex_edgenum, flex_vertadr, flex_vertnum;
 
   // disabled or no equality constraints: return
   if (mjDISABLED(mjDSBL_EQUALITY) || m->nemax == 0) {
@@ -1728,14 +1770,30 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
       }
       break;
 
+    case mjEQ_FLEXVERT:
+      flex_vertadr = m->flex_vertadr[id[0]];
+      flex_vertnum = m->flex_vertnum[id[0]];
+      size = 2 * flex_vertnum;
+      if (nnz) {
+        for (int v=flex_vertadr; v < flex_vertadr+flex_vertnum; v++) {
+          NV += m->flexvert_J_rownnz[2*v+0];
+          NV += m->flexvert_J_rownnz[2*v+1];
+        }
+      }
+      break;
+
     default:
       // might occur in case of the now-removed distance equality constraint
-      mjERROR("unknown constraint type type %d", m->eq_type[i]);    // SHOULD NOT OCCUR
+      mjERROR("unknown constraint type %d", m->eq_type[i]);    // SHOULD NOT OCCUR
     }
 
     // accumulate counts; flex NV already accumulated
     ne += mj_addConstraintCount(m, size, NV);
-    nnze += (m->eq_type[i] == mjEQ_FLEX) ? NV : size*NV;
+    if (m->eq_type[i] == mjEQ_FLEX || m->eq_type[i] == mjEQ_FLEXVERT) {
+      nnze += NV;
+    } else {
+      nnze += size*NV;
+    }
   }
 
   if (nnz) {

src/engine/engine_core_smooth.c

@@ -2775,6 +2775,11 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
       }
       break;
 
+    case mjEQ_FLEXVERT:
+      k = m->eq_obj1id[id];
+      i += 2*m->flex_vertnum[k];
+      break;
+
     default:
       mjERROR("unknown constraint type type %d", m->eq_type[id]);    // SHOULD NOT OCCUR
     }

src/engine/engine_io.c

@@ -1898,6 +1898,7 @@ const char* mj_validateReferences(const mjModel* m) {
       break;
 
     case mjEQ_FLEX:
+    case mjEQ_FLEXVERT:
       if (obj1id >= m->nflex || obj1id < 0) {
         return "Invalid model: eq_obj1id out of bounds.";
       }

src/engine/engine_island.c

@@ -352,7 +352,9 @@ static int findEdges(const mjModel* m, const mjData* d, int* treenedge, int* edg
     // row i is still in the same constraint: skip it,
     if (efc_type == d->efc_type[i] && efc_id == d->efc_id[i]) {
       // 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)) {
+      if (!(efc_type == mjCNSTR_EQUALITY &&
+            (m->eq_type[efc_id] == mjEQ_FLEX ||
+             m->eq_type[efc_id] == mjEQ_FLEXVERT))) {
         continue;
       }
     }

src/engine/engine_sleep.c

@@ -399,6 +399,7 @@ int mj_wakeEquality(const mjModel* m, mjData* d) {
       mjERROR("tendon equality does not yet support sleeping");
       continue;
     case mjEQ_FLEX:
+    case mjEQ_FLEXVERT:
       mjERROR("flex equality does not yet support sleeping");
       continue;
     default:
@@ -641,6 +642,7 @@ static mjtSleepState mj_equalitySleepState(const mjModel* m, const mjData* d, in
       objtype = mjOBJ_TENDON;
       break;
     case mjEQ_FLEX:
+    case mjEQ_FLEXVERT:
       objtype = mjOBJ_FLEX;
       break;
     default: