Add vertex-based flex constraints (dim=2).

PiperOrigin-RevId: 859552050
Change-Id: I61d4b9dc40f041c5b930e5f8810a803e8bccb87a

src/user/user_mesh.cc

@@ -4392,10 +4392,16 @@ void mjCFlex::Compile(const mjVFS* vfs) {
   // compute bounding box coordinates
   vert0_.assign(3*nvert, 0);
   const mjtNum* bvh = tree.Bvh().data();
+  size[0] = bvh[3] - radius;
+  size[1] = bvh[4] - radius;
+  size[2] = bvh[5] - radius;
   for (int j=0; j < nvert; j++) {
     for (int k=0; k < 3; k++) {
-      double size = 2*(bvh[k+3] - radius);
-      vert0_[3*j+k] = (vertxpos[3*j+k] - bvh[k]) / size + 0.5;
+      if (size[k] > mjMINVAL) {
+        vert0_[3*j+k] = (vertxpos[3*j+k] - bvh[k]) / (2*size[k]) + 0.5;
+      } else {
+        vert0_[3*j+k] = 0.5;
+      }
     }
   }
 

src/user/user_model.cc

@@ -1155,6 +1155,7 @@ void mjCModel::Clear() {
   nflexevpair = 0;
   nflextexcoord = 0;
   nJfe = 0;
+  nJfv = 0;
   nmeshvert = 0;
   nmeshnormal = 0;
   nmeshtexcoord = 0;
@@ -2166,6 +2167,31 @@ void mjCModel::SetSizes() {
         nJfe += b->dofnum;
       }
     }
+
+    // compute nJfv
+    std::vector<std::vector<int>> adj(flexes_[i]->nvert);
+    for (const auto& edge : flexes_[i]->edge) {
+      adj[edge.first].push_back(edge.second);
+      adj[edge.second].push_back(edge.first);
+    }
+    for (int j=0; j < flexes_[i]->nvert; j++) {
+      std::unordered_set<int> vert_bodies;
+      vert_bodies.insert(flexes_[i]->vertbodyid[j]);
+      for (int neighbor : adj[j]) {
+        vert_bodies.insert(flexes_[i]->vertbodyid[neighbor]);
+      }
+      std::unordered_set<mjCBody*> bodies_in_jac;
+      for (int body_id : vert_bodies) {
+        mjCBody* b = bodies_[body_id];
+        while (b) {
+          bodies_in_jac.insert(b);
+          b = b->parent;
+        }
+      }
+      for (mjCBody* b : bodies_in_jac) {
+        nJfv += b->dofnum;
+      }
+    }
   }
 
   // mesh counts
@@ -3288,6 +3314,7 @@ void mjCModel::CopyObjects(mjModel* m) {
     mjuu_copyvec(m->flex_solref + mjNREF * i, pfl->solref, mjNREF);
     mjuu_copyvec(m->flex_solimp + mjNIMP * i, pfl->solimp, mjNIMP);
     m->flex_radius[i] = (mjtNum)pfl->radius;
+    mjuu_copyvec(m->flex_size + 3 * i, pfl->size, 3);
     mjuu_copyvec(m->flex_friction + 3 * i, pfl->friction, 3);
     m->flex_margin[i] = (mjtNum)pfl->margin;
     m->flex_gap[i] = (mjtNum)pfl->gap;
@@ -3363,6 +3390,7 @@ void mjCModel::CopyObjects(mjModel* m) {
         m->flex_edgeequality[i] = 1;
         break;
       }
+      // TODO: support flex_edgeequality = 2
     }
 
     // copy bvh data (flex aabb computed dynamically in mjData)
@@ -4856,7 +4884,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
   mj_makeModel(&m,
                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, nJfe,
+               nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord, nJfe, nJfv,
                nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph, nmeshpoly,
                nmeshpolyvert, nmeshpolymap, nskin, nskinvert, nskintexvert, nskinface, nskinbone,
                nskinbonevert, nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude,

src/user/user_model.h

@@ -99,7 +99,8 @@ class mjCModel_ : public mjsElement {
   int nflexshelldata;  // number of shell fragment vertex ids in all flexes
   int nflexevpair;     // number of element-vertex pairs in all flexes
   int nflextexcoord;   // number of vertex texture coordinates in all flexes
-  int nJfe;            // number of non-zeros in sparse flex constraint Jacobian
+  int nJfe;            // number of non-zeros in sparse flex edge constraint Jacobian
+  int nJfv;            // number of non-zeros in sparse flex vertex constraint Jacobian
   int nmeshvert;       // number of vertices in all meshes
   int nmeshnormal;     // number of normals in all meshes
   int nmeshtexcoord;   // number of texture coordinates in all meshes