Add edge ids to flex elements in mjModel.

PiperOrigin-RevId: 675130981
Change-Id: If2ccfb268142d235273b7a50e15740f6399903b2

doc/includes/references.h

@@ -899,6 +899,7 @@ struct mjModel_ {
   int nflexedge;                  // number of edges in all flexes
   int nflexelem;                  // number of elements in all flexes
   int nflexelemdata;              // number of element vertex ids in all flexes
+  int nflexelemedge;              // number of element edge ids in all flexes
   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 vertices with texture coordinates
@@ -1152,6 +1153,7 @@ struct mjModel_ {
   int*      flex_elemadr;         // first element address                    (nflex x 1)
   int*      flex_elemnum;         // number of elements                       (nflex x 1)
   int*      flex_elemdataadr;     // first element vertex id address          (nflex x 1)
+  int*      flex_elemedgeadr;     // first element edge id address            (nflex x 1)
   int*      flex_shellnum;        // number of shells                         (nflex x 1)
   int*      flex_shelldataadr;    // first shell data address                 (nflex x 1)
   int*      flex_evpairadr;       // first evpair address                     (nflex x 1)
@@ -1160,6 +1162,7 @@ struct mjModel_ {
   int*      flex_vertbodyid;      // vertex body ids                          (nflexvert x 1)
   int*      flex_edge;            // edge vertex ids (2 per edge)             (nflexedge x 2)
   int*      flex_elem;            // element vertex ids (dim+1 per elem)      (nflexelemdata x 1)
+  int*      flex_elemedge;        // element edge ids                         (nflexelemedge x 1)
   int*      flex_elemlayer;       // element distance from surface, 3D only   (nflexelem x 1)
   int*      flex_shell;           // shell fragment vertex ids (dim per frag) (nflexshelldata x 1)
   int*      flex_evpair;          // (element, vertex) collision pairs        (nflexevpair x 2)

include/mujoco/mjmodel.h

@@ -609,6 +609,7 @@ struct mjModel_ {
   int nflexedge;                  // number of edges in all flexes
   int nflexelem;                  // number of elements in all flexes
   int nflexelemdata;              // number of element vertex ids in all flexes
+  int nflexelemedge;              // number of element edge ids in all flexes
   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 vertices with texture coordinates
@@ -862,6 +863,7 @@ struct mjModel_ {
   int*      flex_elemadr;         // first element address                    (nflex x 1)
   int*      flex_elemnum;         // number of elements                       (nflex x 1)
   int*      flex_elemdataadr;     // first element vertex id address          (nflex x 1)
+  int*      flex_elemedgeadr;     // first element edge id address            (nflex x 1)
   int*      flex_shellnum;        // number of shells                         (nflex x 1)
   int*      flex_shelldataadr;    // first shell data address                 (nflex x 1)
   int*      flex_evpairadr;       // first evpair address                     (nflex x 1)
@@ -870,6 +872,7 @@ struct mjModel_ {
   int*      flex_vertbodyid;      // vertex body ids                          (nflexvert x 1)
   int*      flex_edge;            // edge vertex ids (2 per edge)             (nflexedge x 2)
   int*      flex_elem;            // element vertex ids (dim+1 per elem)      (nflexelemdata x 1)
+  int*      flex_elemedge;        // element edge ids                         (nflexelemedge x 1)
   int*      flex_elemlayer;       // element distance from surface, 3D only   (nflexelem x 1)
   int*      flex_shell;           // shell fragment vertex ids (dim per frag) (nflexshelldata x 1)
   int*      flex_evpair;          // (element, vertex) collision pairs        (nflexevpair x 2)

include/mujoco/mjxmacro.h

@@ -85,6 +85,7 @@
     X   ( nflexedge )          \
     X   ( nflexelem )          \
     X   ( nflexelemdata )      \
+    X   ( nflexelemedge )      \
     X   ( nflexshelldata )     \
     X   ( nflexevpair )        \
     XMJV( nflextexcoord )      \
@@ -331,6 +332,7 @@
     XMJV( int,     flex_elemadr,          nflex,         1                    ) \
     XMJV( int,     flex_elemnum,          nflex,         1                    ) \
     XMJV( int,     flex_elemdataadr,      nflex,         1                    ) \
+    X   ( int,     flex_elemedgeadr,      nflex,         1                    ) \
     XMJV( int,     flex_shellnum,         nflex,         1                    ) \
     XMJV( int,     flex_shelldataadr,     nflex,         1                    ) \
     X   ( int,     flex_evpairadr,        nflex,         1                    ) \
@@ -339,6 +341,7 @@
     X   ( int,     flex_vertbodyid,       nflexvert,     1                    ) \
     X   ( int,     flex_edge,             nflexedge,     2                    ) \
     XMJV( int,     flex_elem,             nflexelemdata, 1                    ) \
+    X   ( int,     flex_elemedge,         nflexelemedge, 1                    ) \
     XMJV( int,     flex_elemlayer,        nflexelem,     1                    ) \
     XMJV( int,     flex_shell,            nflexshelldata,1                    ) \
     X   ( int,     flex_evpair,           nflexevpair,   2                    ) \

introspect/structs.py

@@ -918,6 +918,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of element vertex ids in all flexes',
              ),
+             StructFieldDecl(
+                 name='nflexelemedge',
+                 type=ValueType(name='int'),
+                 doc='number of element edge ids in all flexes',
+             ),
              StructFieldDecl(
                  name='nflexshelldata',
                  type=ValueType(name='int'),
@@ -2337,6 +2342,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='first element vertex id address          (nflex x 1)',
              ),
+             StructFieldDecl(
+                 name='flex_elemedgeadr',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='first element edge id address            (nflex x 1)',
+             ),
              StructFieldDecl(
                  name='flex_shellnum',
                  type=PointerType(
@@ -2393,6 +2405,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='element vertex ids (dim+1 per elem)      (nflexelemdata x 1)',  # pylint: disable=line-too-long
              ),
+             StructFieldDecl(
+                 name='flex_elemedge',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='element edge ids                         (nflexelemedge x 1)',  # pylint: disable=line-too-long
+             ),
              StructFieldDecl(
                  name='flex_elemlayer',
                  type=PointerType(

plugin/elasticity/elasticity.cc

@@ -19,72 +19,11 @@
 #include <cstdlib>
 #include <sstream>
 #include <string>
-#include <utility>
 #include <vector>
-#include <unordered_map>
 #include <mujoco/mujoco.h>
 
 namespace mujoco::plugin::elasticity {
 
-template <typename T>
-int CreateStencils(std::vector<T>& elements,
-                   std::vector<std::pair<int, int>>& edges,
-                   const std::vector<int>& simplex,
-                   const std::vector<int>& edgeidx) {
-  int ne = 0;
-  int nt = simplex.size() / T::kNumVerts;
-  elements.resize(nt);
-  for (int t = 0; t < nt; t++) {
-    for (int v = 0; v < T::kNumVerts; v++) {
-      elements[t].vertices[v] = simplex[T::kNumVerts*t+v];
-    }
-  }
-
-  // map from edge vertices to their index in `edges` vector
-  std::unordered_map<std::pair<int, int>, int, PairHash> edge_indices;
-
-  // loop over all tetrahedra
-  for (int t = 0; t < nt; t++) {
-    int* v = elements[t].vertices;
-
-    // compute edges to vertices map for fast computations
-    for (int e = 0; e < T::kNumEdges; e++) {
-      auto pair = std::pair(
-        std::min(v[T::edge[e][0]], v[T::edge[e][1]]),
-        std::max(v[T::edge[e][0]], v[T::edge[e][1]])
-      );
-
-      // if edge is already present in the vector only store its index
-      auto [it, inserted] = edge_indices.insert({pair, ne});
-
-      if (inserted) {
-        edges.push_back(pair);
-        elements[t].edges[e] = ne++;
-      } else {
-        elements[t].edges[e] = it->second;
-      }
-
-      if (!edgeidx.empty()) {  // SHOULD NOT OCCUR
-        if (elements[t].edges[e] != edgeidx[T::kNumEdges*t+e]) {
-          mju_error("edge ordering is incoherent between flex and plugin");
-        }
-      }
-    }
-  }
-
-  return nt;
-}
-
-template int CreateStencils<Stencil2D>(std::vector<Stencil2D>& elements,
-                                       std::vector<std::pair<int, int>>& edges,
-                                       const std::vector<int>& simplex,
-                                       const std::vector<int>& edgeidx);
-
-template int CreateStencils<Stencil3D>(std::vector<Stencil3D>& elements,
-                                       std::vector<std::pair<int, int>>& edges,
-                                       const std::vector<int>& simplex,
-                                       const std::vector<int>& edgeidx);
-
 void String2Vector(const std::string& txt, std::vector<int>& vec) {
   std::stringstream strm(txt);
   vec.clear();

plugin/elasticity/elasticity.h

@@ -65,20 +65,19 @@ void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3],
 
 template <typename T>
 inline void ComputeForce(std::vector<mjtNum>& qfrc_passive,
-                         const std::vector<T>& elements,
                          const std::vector<mjtNum>& elongationglob,
                          const mjModel* m, int flex,
                          const mjtNum* xpos) {
   mju_zero(qfrc_passive.data(), qfrc_passive.size());
   mjtNum* k = m->flex_stiffness + 21 * m->flex_elemadr[flex];
 
-  if (elements.size() != m->flex_elemnum[flex]) {
+  int dim = m->flex_dim[flex];
-    mju_error("plugin stencil does not match flex stencil");
+  const int* elem = m->flex_elem + m->flex_elemdataadr[flex];
-  }
+  const int* edgeelem = m->flex_elemedge + m->flex_elemedgeadr[flex];
 
   // compute force element-by-element
   for (int t = 0; t < m->flex_elemnum[flex]; t++)  {
-    const int* v = elements[t].vertices;
+    const int* v = elem + (dim+1) * t;
 
     // compute length gradient with respect to dofs
     mjtNum gradient[T::kNumEdges][2][3];
@@ -87,7 +86,7 @@ inline void ComputeForce(std::vector<mjtNum>& qfrc_passive,
     // extract elongation of edges belonging to this element
     mjtNum elongation[T::kNumEdges];
     for (int e = 0; e < T::kNumEdges; e++) {
-      int idx = elements[t].edges[e];
+      int idx = edgeelem[t * T::kNumEdges + e];
       elongation[e] = elongationglob[idx];
     }
 

plugin/elasticity/membrane.cc

@@ -87,10 +87,7 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
         strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
     mjtNum damp =
             strtod(mj_getPluginConfig(m, instance, "damping"), nullptr);
-    std::vector<int> face, edge;
+    return Membrane(m, d, instance, nu, E, thick, damp);
-    String2Vector(mj_getPluginConfig(m, instance, "face"), face);
-    String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
-    return Membrane(m, d, instance, nu, E, thick, damp, face, edge);
   } else {
     mju_warning("Invalid parameter specification in shell plugin");
     return std::nullopt;
@@ -99,9 +96,7 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
 
 // plugin constructor
 Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
-                   mjtNum E, mjtNum thick, mjtNum damp,
+                   mjtNum E, mjtNum thick, mjtNum damp)
-                   const std::vector<int>& simplex,
-                   const std::vector<int>& edgeidx)
     : f0(-1), damping(damp), thickness(thick) {
   // count plugin bodies
   nv = ne = 0;
@@ -126,12 +121,10 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   // vertex positions
   mjtNum* body_pos = m->flex_xvert0 + 3*m->flex_vertadr[f0];
 
-  // generate triangles from the vertices
-  nt = CreateStencils<Stencil2D>(elements, edges, simplex, edgeidx);
-
   // loop over all triangles
-  for (int t = 0; t < nt; t++) {
+  const int* elem = m->flex_elem + m->flex_elemdataadr[f0];
-    int* v = elements[t].vertices;
+  for (int t = 0; t < m->flex_elemnum[f0]; t++) {
+    const int* v = elem + (m->flex_dim[f0]+1) * t;
     for (int i = 0; i < kNumVerts; i++) {
       int bi = m->flex_vertbodyid[m->flex_vertadr[f0]+v[i]];
       if (bi && m->body_plugin[bi] != instance) {
@@ -163,7 +156,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   }
 
   // allocate array
-  ne = edges.size();
+  ne = m->flex_edgenum[f0];
   elongation.assign(ne, 0);
   force.assign(3*nv, 0);
 }
@@ -195,7 +188,7 @@ void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
   mjtNum* xpos = d->flexvert_xpos + 3*flex_vertadr;
   mjtNum* qfrc = d->qfrc_passive;
 
-  ComputeForce<Stencil2D>(force, elements, elongation, m, f0, xpos);
+  ComputeForce<Stencil2D>(force, elongation, m, f0, xpos);
 
   // insert into passive force
   AddFlexForce(qfrc, force, m, d, xpos, f0);

plugin/elasticity/membrane.h

@@ -44,13 +44,8 @@ class Membrane {
   int i0;  // index of first body
   int nc;  // number of quads in the grid
   int nv;  // number of vertices (bodies) in the Membrane
-  int nt;  // number of area elements (triangles)
   int ne;  // number of edges in the Membrane
 
-  // connectivity info for mapping tetrahedra to edges and vertices
-  std::vector<Stencil2D> elements;          // triangles              (nt x 6)
-  std::vector<std::pair<int, int> > edges;  // edge to vertex map     (ne x 2)
-
   // precomputed quantities
   std::vector<mjtNum> prev;                 // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
@@ -61,8 +56,7 @@ class Membrane {
 
  private:
   Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-           mjtNum thick, mjtNum damp, const std::vector<int>& simplex,
+           mjtNum thick, mjtNum damp);
-           const std::vector<int>& edgeidx);
 };
 
 }  // namespace mujoco::plugin::elasticity

plugin/elasticity/solid.cc

@@ -93,10 +93,7 @@ std::optional<Solid> Solid::Create(const mjModel* m, mjData* d, int instance) {
         mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
         mjtNum damp =
             strtod(mj_getPluginConfig(m, instance, "damping"), nullptr);
-        std::vector<int> face, edge;
+        return Solid(m, d, instance, nu, E, damp);
-        String2Vector(mj_getPluginConfig(m, instance, "face"), face);
-        String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
-        return Solid(m, d, instance, nu, E, damp, face, edge);
     } else {
         mju_warning("Invalid parameter specification in solid plugin");
         return std::nullopt;
@@ -105,8 +102,7 @@ std::optional<Solid> Solid::Create(const mjModel* m, mjData* d, int instance) {
 
 // plugin constructor
 Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-             mjtNum damp, const std::vector<int>& simplex,
+             mjtNum damp)
-             const std::vector<int>& edgeidx)
     : f0(-1), damping(damp) {
   // count plugin bodies
   nv = ne = 0;
@@ -134,12 +130,10 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   // vertex positions
   mjtNum* body_pos = m->flex_xvert0 + 3*m->flex_vertadr[f0];
 
-  // generate tetrahedra from the vertices
-  nt = CreateStencils<Stencil3D>(elements, edges, simplex, edgeidx);
-
   // loop over all tetrahedra
-  for (int t = 0; t < nt; t++) {
+  const int* elem = m->flex_elem + m->flex_elemdataadr[f0];
-    int* v = elements[t].vertices;
+  for (int t = 0; t < m->flex_elemnum[f0]; t++) {
+    const int* v = elem + (m->flex_dim[f0]+1) * t;
     for (int i = 0; i < kNumVerts; i++) {
       int bi = m->flex_vertbodyid[m->flex_vertadr[f0]+v[i]];
       if (bi && m->body_plugin[bi] != instance) {
@@ -170,7 +164,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   }
 
   // allocate array
-  ne = edges.size();
+  ne = m->flex_edgenum[f0];
   elongation.assign(ne, 0);
   force.assign(3*nv, 0);
 }
@@ -202,7 +196,7 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) {
   mjtNum* xpos = d->flexvert_xpos + 3*flex_vertadr;
   mjtNum* qfrc = d->qfrc_passive;
 
-  ComputeForce<Stencil3D>(force, elements, elongation, m, f0, xpos);
+  ComputeForce<Stencil3D>(force, elongation, m, f0, xpos);
 
   // insert into passive force
   AddFlexForce(qfrc, force, m, d, xpos, f0);

plugin/elasticity/solid.h

@@ -42,13 +42,8 @@ class Solid {
   int i0;  // index of first body
   int nc;  // number of cubes in the grid
   int nv;  // number of vertices (bodies) in the solid
-  int nt;  // number of volumetric elements (tetrahedra)
   int ne;  // number of edges in the solid
 
-  // connectivity info for mapping tetrahedra to edges and vertices
-  std::vector<Stencil3D> elements;          // 4 vertices and 6 edges (nt x 10)
-  std::vector<std::pair<int, int> > edges;  // edge to vertex map     (ne x 2)
-
   // precomputed quantities
   std::vector<mjtNum> prev;                 // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
@@ -58,8 +53,7 @@ class Solid {
 
  private:
   Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-        mjtNum damp, const std::vector<int>& simplex,
+        mjtNum damp);
-        const std::vector<int>& edgeidx);
 };
 
 }  // namespace mujoco::plugin::elasticity

src/engine/engine_io.c

@@ -461,7 +461,7 @@ void mj_makeModel(mjModel** dest,
     int nq, int nv, int nu, int na, int nbody, int nbvh,
     int nbvhstatic, int nbvhdynamic, int njnt, int ngeom, int nsite, int ncam,
     int nlight, int nflex, int nflexvert, int nflexedge, int nflexelem,
-    int nflexelemdata, int nflexshelldata, int nflexevpair, int nflextexcoord,
+    int nflexelemdata, int nflexelemedge, int nflexshelldata, int nflexevpair, int nflextexcoord,
     int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
     int nmeshgraph, int nskin, int nskinvert, int nskintexvert, int nskinface,
     int nskinbone, int nskinbonevert, int nhfield, int nhfielddata, int ntex,
@@ -507,6 +507,7 @@ void mj_makeModel(mjModel** dest,
   m->nflexedge = nflexedge;
   m->nflexelem = nflexelem;
   m->nflexelemdata = nflexelemdata;
+  m->nflexelemedge = nflexelemedge;
   m->nflexshelldata = nflexshelldata;
   m->nflexevpair = nflexevpair;
   m->nflextexcoord = nflextexcoord;
@@ -634,7 +635,7 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) {
       src->nq, src->nv, src->nu, src->na, src->nbody, src->nbvh,
       src->nbvhstatic, src->nbvhdynamic, src->njnt, src->ngeom, src->nsite,
       src->ncam, src->nlight, src->nflex, src->nflexvert, src->nflexedge,
-      src->nflexelem, src->nflexelemdata, src->nflexshelldata,
+      src->nflexelem, src->nflexelemdata, src->nflexelemedge, src->nflexshelldata,
       src->nflexevpair, src->nflextexcoord, src->nmesh, src->nmeshvert,
       src->nmeshnormal, src->nmeshtexcoord, src->nmeshface, src->nmeshgraph,
       src->nskin, src->nskinvert, src->nskintexvert, src->nskinface,
@@ -795,7 +796,7 @@ mjModel* mj_loadModelBuffer(const void* buffer, int buffer_sz) {
                ints[42], ints[43], ints[44], ints[45], ints[46], ints[47], ints[48],
                ints[49], ints[50], ints[51], ints[52], ints[53], ints[54], ints[55],
                ints[56], ints[57], ints[58], ints[59], ints[60], ints[61], ints[62],
-               ints[63]);
+               ints[63], ints[64]);
   if (!m || m->nbuffer != sizes[getnsize()-1]) {
     mju_warning("Corrupted model, wrong size parameters");
     mj_deleteModel(m);
@@ -2097,9 +2098,11 @@ const char* mj_validateReferences(const mjModel* m) {
   X(flex_evpairadr,     nflex,          nflexevpair   , m->flex_evpairnum      ) \
   X(flex_texcoordadr,   nflex,          nflextexcoord , 0                      ) \
   X(flex_elemdataadr,   nflex,          nflexelemdata , 0                      ) \
+  X(flex_elemedgeadr,   nflex,          nflexelemedge , 0                      ) \
   X(flex_shelldataadr,  nflex,          nflexshelldata, 0                      ) \
   X(flex_edge,          nflexedge*2,    nflexvert     , 0                      ) \
   X(flex_elem,          nflexelemdata,  nflexvert     , 0                      ) \
+  X(flex_elemedge,      nflexelemedge,  nflexedge     , 0                      ) \
   X(flex_shell,         nflexshelldata, nflexvert     , 0                      ) \
   X(flex_bvhadr,        nflex,          nbvh          , m->flex_bvhnum         ) \
   X(skin_matid,         nskin,          nmat          , 0                      ) \

src/engine/engine_io.h

@@ -54,10 +54,10 @@ void mj_defaultStatistic(mjStatistic* stat);
 void mj_makeModel(mjModel** dest,
     int nq, int nv, int nu, int na, int nbody, int nbvh, int nbvhstatic, int nbvhdynamic,
     int njnt, int ngeom, int nsite, int ncam, int nlight, int nflex, int nflexvert,
-    int nflexedge, int nflexelem, int nflexelemdata, int nflexshelldata, int nflexevpair,
+    int nflexedge, int nflexelem, int nflexelemdata, int nflexelemedge, int nflexshelldata,
-    int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal, int nmeshtexcoord, int nmeshface,
+    int nflexevpair, int nflextexcoord, int nmesh, int nmeshvert, int nmeshnormal,
-    int nmeshgraph, int nskin, int nskinvert, int nskintexvert, int nskinface,
+    int nmeshtexcoord, int nmeshface, int nmeshgraph, int nskin, int nskinvert, int nskintexvert,
-    int nskinbone, int nskinbonevert, int nhfield, int nhfielddata,
+    int nskinface, int nskinbone, int nskinbonevert, int nhfield, int nhfielddata,
     int ntex, int ntexdata, int nmat, int npair, int nexclude,
     int neq, int ntendon, int nwrap, int nsensor,
     int nnumeric, int nnumericdata, int ntext, int ntextdata,

src/user/user_mesh.cc

@@ -2386,7 +2386,6 @@ struct PairHash
 };
 
 // simplex connectivity
-constexpr int kNumEdges[3] = {1, 3, 6};
 constexpr int eledge[3][6][2] = {{{ 0,  1}, {-1, -1}, {-1, -1},
                                   {-1, -1}, {-1, -1}, {-1, -1}},
                                  {{ 1,  2}, { 2,  0}, { 0,  1},
@@ -2637,7 +2636,7 @@ void mjCFlex::Compile(const mjVFS* vfs) {
   }
 
   // create edges
-  std::vector<int> edgeidx(elem_.size()*kNumEdges[dim-1]);
+  edgeidx_.assign(elem_.size()*kNumEdges[dim-1]/(dim+1), 0);
 
   // map from edge vertices to their index in `edges` vector
   std::unordered_map<std::pair<int, int>, int, PairHash> edge_indices;
@@ -2656,9 +2655,9 @@ void mjCFlex::Compile(const mjVFS* vfs) {
 
       if (inserted) {
         edge.push_back(pair);
-        edgeidx[f*kNumEdges[dim-1]+e] = nedge++;
+        edgeidx_[f*kNumEdges[dim-1]+e] = nedge++;
       } else {
-        edgeidx[f*kNumEdges[dim-1]+e] = it->second;
+        edgeidx_[f*kNumEdges[dim-1]+e] = it->second;
       }
     }
   }
@@ -2669,7 +2668,7 @@ void mjCFlex::Compile(const mjVFS* vfs) {
   // add plugins
   std::string userface, useredge;
   userface = VectorToString(elem_);
-  useredge = VectorToString(edgeidx);
+  useredge = VectorToString(edgeidx_);
 
   for (const auto& vbodyid : vertbodyid) {
     if (model->Bodies()[vbodyid]->plugin.instance) {

src/user/user_model.cc

@@ -654,6 +654,7 @@ void mjCModel::Clear() {
   nflexedge = 0;
   nflexelem = 0;
   nflexelemdata = 0;
+  nflexelemedge = 0;
   nflexshelldata = 0;
   nflexevpair = 0;
   nflextexcoord = 0;
@@ -1545,6 +1546,7 @@ void mjCModel::SetSizes() {
     nflexedge += flexes_[i]->nedge;
     nflexelem += flexes_[i]->nelem;
     nflexelemdata += flexes_[i]->nelem * (flexes_[i]->dim + 1);
+    nflexelemedge += flexes_[i]->nelem * mjCFlex::kNumEdges[flexes_[i]->dim - 1];
     nflexshelldata += (int)flexes_[i]->shell.size();
     nflexevpair += (int)flexes_[i]->evpair.size()/2;
   }
@@ -2409,7 +2411,7 @@ void mjCModel::CopyTree(mjModel* m) {
 // copy objects outside kinematic tree
 void mjCModel::CopyObjects(mjModel* m) {
   int adr, bone_adr, vert_adr, normal_adr, face_adr, texcoord_adr;
-  int edge_adr, elem_adr, elemdata_adr, shelldata_adr, evpair_adr;
+  int edge_adr, elem_adr, elemdata_adr, elemedge_adr, shelldata_adr, evpair_adr;
   int bonevert_adr, graph_adr, data_adr, bvh_adr;
 
   // sizes outside call to mj_makeModel
@@ -2491,6 +2493,7 @@ void mjCModel::CopyObjects(mjModel* m) {
   edge_adr = 0;
   elem_adr = 0;
   elemdata_adr = 0;
+  elemedge_adr = 0;
   shelldata_adr = 0;
   evpair_adr = 0;
   texcoord_adr = 0;
@@ -2526,6 +2529,7 @@ void mjCModel::CopyObjects(mjModel* m) {
     m->flex_edgenum[i] = pfl->nedge;
     m->flex_elemadr[i] = elem_adr;
     m->flex_elemdataadr[i] = elemdata_adr;
+    m->flex_elemedgeadr[i] = elemedge_adr;
     m->flex_shellnum[i] = (int)pfl->shell.size()/pfl->dim;
     m->flex_shelldataadr[i] = m->flex_shellnum[i] ? shelldata_adr : -1;
     if (pfl->evpair.empty()) {
@@ -2545,6 +2549,7 @@ void mjCModel::CopyObjects(mjModel* m) {
     }
     m->flex_elemnum[i] = pfl->nelem;
     memcpy(m->flex_elem + elemdata_adr, pfl->elem_.data(), pfl->elem_.size()*sizeof(int));
+    memcpy(m->flex_elemedge + elemedge_adr, pfl->edgeidx_.data(), pfl->edgeidx_.size()*sizeof(int));
     memcpy(m->flex_elemlayer + elem_adr, pfl->elemlayer.data(), pfl->nelem*sizeof(int));
     if (m->flex_shellnum[i]) {
       memcpy(m->flex_shell + shelldata_adr, pfl->shell.data(), pfl->shell.size()*sizeof(int));
@@ -2616,6 +2621,7 @@ void mjCModel::CopyObjects(mjModel* m) {
     edge_adr += pfl->nedge;
     elem_adr += pfl->nelem;
     elemdata_adr += (pfl->dim+1) * pfl->nelem;
+    elemedge_adr += (pfl->kNumEdges[pfl->dim-1]) * pfl->nelem;
     shelldata_adr += (int)pfl->shell.size();
     evpair_adr += (int)pfl->evpair.size()/2;
     texcoord_adr += (int)pfl->texcoord_.size()/2;
@@ -3897,7 +3903,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
   mj_makeModel(&m,
                nq, nv, nu, na, nbody, nbvh, nbvhstatic, nbvhdynamic, njnt, ngeom, nsite,
                ncam, nlight, nflex, nflexvert, nflexedge, nflexelem,
-               nflexelemdata, nflexshelldata, nflexevpair, nflextexcoord,
+               nflexelemdata, nflexelemedge, nflexshelldata, nflexevpair, nflextexcoord,
                nmesh, nmeshvert, nmeshnormal, nmeshtexcoord, nmeshface, nmeshgraph,
                nskin, nskinvert, nskintexvert, nskinface, nskinbone, nskinbonevert,
                nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude,

src/user/user_model.h

@@ -90,6 +90,7 @@ class mjCModel_ : public mjsElement {
   int nflexedge;       // number of edges in all flexes
   int nflexelem;       // number of elements in all flexes
   int nflexelemdata;   // number of element vertex ids in all flexes
+  int nflexelemedge;   // number of element edges in all flexes
   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

src/user/user_objects.h

@@ -694,6 +694,7 @@ class mjCFlex_ : public mjCBase {
   std::vector<double> vertxpos;           // global vertex positions
   mjCBoundingVolumeHierarchy tree;        // bounding volume hierarchy
   std::vector<double> elemaabb_;          // element bounding volume
+  std::vector<int> edgeidx_;              // element edge ids
 
   // variable-size data
   std::vector<std::string> vertbody_;     // vertex body names
@@ -741,6 +742,8 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
   bool HasTexcoord() const;               // texcoord not null
   void DelTexcoord();                     // delete texcoord
 
+  static constexpr int kNumEdges[3] = {1, 3, 6};  // number of edges per element indexed by dim
+
  private:
   void Compile(const mjVFS* vfs);         // compiler
   void CreateBVH(void);                   // create flex BVH

test/plugin/elasticity/elasticity_test.cc

@@ -138,8 +138,6 @@ TEST_F(PluginTest, ElasticEnergyMembrane) {
   mjModel* m = LoadModelFromString(cantilever_xml, error, sizeof(error));
   ASSERT_THAT(m, testing::NotNull()) << error;
   mjData* d = mj_makeData(m);
-  auto* membrane =
-      reinterpret_cast<plugin::elasticity::Membrane*>(d->plugin_data[0]);
 
   mj_kinematics(m, d);
   mj_flex(m, d);
@@ -149,14 +147,14 @@ TEST_F(PluginTest, ElasticEnergyMembrane) {
   // trace(strain^2) = 2*scale^2
 
   for (mjtNum scale = 1; scale < 4; scale++) {
-    for (int t = 0; t < membrane->nt; t++) {
+    for (int t = 0; t < m->flex_elemnum[0]; t++) {
       mjtNum energy = 0;
       mjtNum volume = 1./2.;
       int idx = 0;
       for (int e1 = 0; e1 < 3; e1++) {
         for (int e2 = e1; e2 < 3; e2++) {
-          int idx1 = membrane->elements[t].edges[e1] + m->flex_edgeadr[0];
+          int idx1 = m->flex_elemedge[3*t+e1 + m->flex_elemedgeadr[0]];
-          int idx2 = membrane->elements[t].edges[e2] + m->flex_edgeadr[0];
+          int idx2 = m->flex_elemedge[3*t+e2 + m->flex_elemedgeadr[0]];
           mjtNum elong1 =
               scale * m->flexedge_length0[idx1] * m->flexedge_length0[idx1];
           mjtNum elong2 =
@@ -221,7 +219,6 @@ TEST_F(ElasticityTest, ElasticEnergySolid) {
   mjModel* m = LoadModelFromString(cantilever_xml, error, sizeof(error));
   ASSERT_THAT(m, testing::NotNull()) << error;
   mjData* d = mj_makeData(m);
-  auto* solid = reinterpret_cast<plugin::elasticity::Solid*>(d->plugin_data[0]);
 
   mj_kinematics(m, d);
   mj_flex(m, d);
@@ -231,14 +228,14 @@ TEST_F(ElasticityTest, ElasticEnergySolid) {
   // trace(strain^2) = 3*scale^2
 
   for (mjtNum scale = 1; scale < 4; scale++) {
-    for (int t = 0; t < solid->nt; t++) {
+    for (int t = 0; t < m->flex_elemnum[0]; t++) {
       mjtNum energy = 0;
       mjtNum volume = 1./6.;
       int idx = 0;
       for (int e1 = 0; e1 < 6; e1++) {
         for (int e2 = e1; e2 < 6; e2++) {
-          int idx1 = solid->elements[t].edges[e1] + m->flex_edgeadr[0];
+          int idx1 = m->flex_elemedge[6*t+e1 + m->flex_elemedgeadr[0]];
-          int idx2 = solid->elements[t].edges[e2] + m->flex_edgeadr[0];
+          int idx2 = m->flex_elemedge[6*t+e2 + m->flex_elemedgeadr[0]];
           mjtNum elong1 =
               scale * m->flexedge_length0[idx1] * m->flexedge_length0[idx1];
           mjtNum elong2 =

unity/Runtime/Bindings/MjBindings.cs

@@ -5184,6 +5184,7 @@ public unsafe struct mjModel_ {
   public int nflexedge;
   public int nflexelem;
   public int nflexelemdata;
+  public int nflexelemedge;
   public int nflexshelldata;
   public int nflexevpair;
   public int nflextexcoord;
@@ -5405,6 +5406,7 @@ public unsafe struct mjModel_ {
   public int* flex_elemadr;
   public int* flex_elemnum;
   public int* flex_elemdataadr;
+  public int* flex_elemedgeadr;
   public int* flex_shellnum;
   public int* flex_shelldataadr;
   public int* flex_evpairadr;
@@ -5413,6 +5415,7 @@ public unsafe struct mjModel_ {
   public int* flex_vertbodyid;
   public int* flex_edge;
   public int* flex_elem;
+  public int* flex_elemedge;
   public int* flex_elemlayer;
   public int* flex_shell;
   public int* flex_evpair;