Fix flex constraint Jacobian for edges belonging to a single body.

PiperOrigin-RevId: 575161094
Change-Id: If4153bcb751cdd6dbdddeb08672f7107c7c5bbc2

doc/includes/references.h

@@ -944,7 +944,7 @@ struct mjModel_ {
   int*      body_treeid;          // id of body's kinematic tree; -1: static  (nbody x 1)
   int*      body_geomnum;         // number of geoms                          (nbody x 1)
   int*      body_geomadr;         // start addr of geoms; -1: no geoms        (nbody x 1)
-  mjtByte*  body_simple;          // 1: diagonal M; 2: diag M, no rotations   (nbody x 1)
+  mjtByte*  body_simple;          // 1: diag M; 2: diag M, sliders only       (nbody x 1)
   mjtByte*  body_sameframe;       // inertial frame is same as body frame     (nbody x 1)
   mjtNum*   body_pos;             // position offset rel. to parent body      (nbody x 3)
   mjtNum*   body_quat;            // orientation offset rel. to parent body   (nbody x 4)
@@ -1122,6 +1122,7 @@ struct mjModel_ {
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
   mjtByte*  flex_edgeequality;    // is edge equality constraint defined      (nflex x 1)
   mjtByte*  flex_rigid;           // are all verices in the same body         (nflex x 1)
+  mjtByte*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
   mjtByte*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
   mjtByte*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
   int*      flex_bvhadr;          // address of bvh root; -1: no bvh          (nflex x 1)

include/mujoco/mjmodel.h

@@ -661,7 +661,7 @@ struct mjModel_ {
   int*      body_treeid;          // id of body's kinematic tree; -1: static  (nbody x 1)
   int*      body_geomnum;         // number of geoms                          (nbody x 1)
   int*      body_geomadr;         // start addr of geoms; -1: no geoms        (nbody x 1)
-  mjtByte*  body_simple;          // 1: diagonal M; 2: diag M, no rotations   (nbody x 1)
+  mjtByte*  body_simple;          // 1: diag M; 2: diag M, sliders only       (nbody x 1)
   mjtByte*  body_sameframe;       // inertial frame is same as body frame     (nbody x 1)
   mjtNum*   body_pos;             // position offset rel. to parent body      (nbody x 3)
   mjtNum*   body_quat;            // orientation offset rel. to parent body   (nbody x 4)
@@ -839,6 +839,7 @@ struct mjModel_ {
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
   mjtByte*  flex_edgeequality;    // is edge equality constraint defined      (nflex x 1)
   mjtByte*  flex_rigid;           // are all verices in the same body         (nflex x 1)
+  mjtByte*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
   mjtByte*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
   mjtByte*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
   int*      flex_bvhadr;          // address of bvh root; -1: no bvh          (nflex x 1)

include/mujoco/mjxmacro.h

@@ -345,6 +345,7 @@
     X   ( mjtNum,  flex_edgedamping,      nflex,         1                    ) \
     X   ( mjtByte, flex_edgeequality,     nflex,         1                    ) \
     X   ( mjtByte, flex_rigid,            nflex,         1                    ) \
+    X   ( mjtByte, flexedge_rigid,        nflexedge,     1                    ) \
     X   ( mjtByte, flex_centered,         nflex,         1                    ) \
     XMJV( mjtByte, flex_flatskin,         nflex,         1                    ) \
     XMJV( int,     flex_bvhadr,           nflex,         1                    ) \

introspect/structs.py

@@ -1311,7 +1311,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=PointerType(
                      inner_type=ValueType(name='mjtByte'),
                  ),
-                 doc='1: diagonal M; 2: diag M, no rotations   (nbody x 1)',
+                 doc='1: diag M; 2: diag M, sliders only       (nbody x 1)',
              ),
              StructFieldDecl(
                  name='body_sameframe',
@@ -2426,6 +2426,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='are all verices in the same body         (nflex x 1)',
              ),
+             StructFieldDecl(
+                 name='flexedge_rigid',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtByte'),
+                 ),
+                 doc='are both edge vertices in same body      (nflexedge x 1)',
+             ),
              StructFieldDecl(
                  name='flex_centered',
                  type=PointerType(

src/engine/engine_core_constraint.c

@@ -663,9 +663,13 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
       case mjEQ_FLEX:
         flex_edgeadr = m->flex_edgeadr[id[0]];
         flex_edgenum = m->flex_edgenum[id[0]];
-
+        // add one constraint per non-rigid edge
-        // add one constraint per edge
         for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
+          // skip rigid
+          if (m->flexedge_rigid[e]) {
+            continue;
+          }
+
           // position error
           cpos[0] = d->flexedge_length[e] - m->flexedge_length0[e];
 
@@ -1057,7 +1061,7 @@ void mj_instantiateContact(const mjModel* m, mjData* d) {
 
 // compute diagApprox
 void mj_diagApprox(const mjModel* m, mjData* d) {
-  int id, dim, b1, b2, weldcnt = 0, edgecnt = 0;
+  int id, dim, b1, b2, f, weldcnt = 0;
   int nefc = d->nefc;
   mjtNum tran, rot, fri, *dA = d->efc_diagApprox;
   mjContact* con = NULL;
@@ -1067,11 +1071,6 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
     // get constraint id
     id = d->efc_id[i];
 
-    // clear edge counter
-    if (d->efc_type[i] != mjEQ_FLEX) {
-      edgecnt = 0;
-    }
-
     // process according to constraint type
     switch ((mjtConstraint) d->efc_type[i]) {
     case mjCNSTR_EQUALITY:
@@ -1108,8 +1107,18 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
         break;
 
       case mjEQ_FLEX:
-        dA[i] = m->flexedge_invweight0[m->flex_edgeadr[m->eq_obj1id[id]] + edgecnt];
+        // process all non-rigid edges for this flex
-        edgecnt++;
+        f = m->eq_obj1id[id];
+        int flex_edgeadr = m->flex_edgeadr[f];
+        int flex_edgenum = m->flex_edgenum[f];
+        for (int e=flex_edgeadr; e<flex_edgeadr+flex_edgenum; e++) {
+          if (!m->flexedge_rigid[e]) {
+            dA[i++] = m->flexedge_invweight0[e];
+          }
+        }
+
+        // adjust constraint counter
+        i--;
         break;
 
       default:
@@ -1642,21 +1651,26 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
         break;
 
       case mjEQ_FLEX:
-        size = m->flex_edgenum[id[0]];
-        if (!nnz) {
-          break;
-        }
-
         flex_edgeadr = m->flex_edgeadr[id[0]];
         flex_edgenum = m->flex_edgenum[id[0]];
 
+        // init with all edges, subract rigid later
+        size = flex_edgenum;
+
         // process edges of this flex
         for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
-          int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]];
+          // rigid: reduce size and skip
-          int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]];
+          if (m->flexedge_rigid[e]) {
+            size--;
+            continue;
+          }
 
-          // accumulate NV
+          // accumulate NV if needed
-          NV += mj_jacDifPairCount(m, chain, b1, b2, issparse);
+          if (nnz) {
+            int b1 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e]];
+            int b2 = m->flex_vertbodyid[m->flex_vertadr[id[0]] + m->flex_edge[2*e+1]];
+            NV += mj_jacDifPairCount(m, chain, b1, b2, issparse);
+          }
         }
         break;
 

src/engine/engine_core_smooth.c

@@ -1857,8 +1857,16 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
       break;
 
     case mjEQ_FLEX:
-      // increment edgenum rows
+      // increment with number of non-rigid edges
-      i += m->flex_edgenum[m->eq_obj1id[id]];
+      k = m->eq_obj1id[id];
+      int flex_edgeadr = m->flex_edgeadr[k];
+      int flex_edgenum = m->flex_edgenum[k];
+
+      for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
+        if (!m->flexedge_rigid[e]) {
+          i++;
+        }
+      }
       break;
 
     default:

src/engine/engine_derivative.c

@@ -1474,9 +1474,16 @@ void mjd_passive_vel(const mjModel* m, mjData* d) {
   for (int f=0; f < m->nflex; f++) {
     if (!m->flex_rigid[f] && m->flex_edgedamping[f]) {
       mjtNum B = -m->flex_edgedamping[f];
+      int flex_edgeadr = m->flex_edgeadr[f];
+      int flex_edgenum = m->flex_edgenum[f];
+
+      // process non-rigid edges of this flex
+      for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
+        // skip rigid
+        if (m->flexedge_rigid[e]) {
+          continue;
+        }
 
-      // process edges of this flex
-      for (int e=m->flex_edgeadr[f]; e < m->flex_edgeadr[f]+m->flex_edgenum[f]; e++) {
         // add sparse or dense
         if (mj_isSparse(m)) {
           addJTBJSparse(m, d, d->flexedge_J, &B, 1, e,

src/engine/engine_passive.c

@@ -106,9 +106,14 @@ void mj_passive(const mjModel* m, mjData* d) {
       continue;
     }
 
-    // process edges of this flex (global edge index)
+    // process non-rigid edges of this flex (global edge index)
     int edgeend = m->flex_edgeadr[f] + m->flex_edgenum[f];
     for (int e=m->flex_edgeadr[f]; e < edgeend; e++) {
+      // skip rigid
+      if (m->flexedge_rigid[e]) {
+        continue;
+      }
+
       // compute spring-damper force along edge
       frc = stiffness * (m->flexedge_length0[e] - d->flexedge_length[e])
             - damping * d->flexedge_velocity[e];

src/engine/engine_sensor.c

@@ -956,10 +956,14 @@ void mj_energyPos(const mjModel* m, mjData* d) {
         continue;
       }
 
-      // process edges of this flex
+      // process non-rigid edges of this flex
-      for (int e=m->flex_edgeadr[i]; e < m->flex_edgeadr[i]+m->flex_edgenum[i]; e++) {
+      int flex_edgeadr = m->flex_edgeadr[i];
-        mjtNum displacement = m->flexedge_length0[e] - d->flexedge_length[e];
+      int flex_edgenum = m->flex_edgenum[i];
-        d->energy[0] += 0.5*stiffness*displacement*displacement;
+      for (int e=flex_edgeadr; e < flex_edgeadr+flex_edgenum; e++) {
+        if (!m->flexedge_rigid[e]) {
+          mjtNum displacement = m->flexedge_length0[e] - d->flexedge_length[e];
+          d->energy[0] += 0.5*stiffness*displacement*displacement;
+        };
       }
     }
   }

src/engine/engine_setconst.c

@@ -206,8 +206,13 @@ static void set0(mjModel* m, mjData* d) {
         int b1 = m->flex_vertbodyid[m->flex_vertadr[f] + m->flex_edge[2*i]];
         int b2 = m->flex_vertbodyid[m->flex_vertadr[f] + m->flex_edge[2*i+1]];
 
+        // rigid edge: set to 0
+        if (m->flexedge_rigid[i]) {
+          m->flexedge_invweight0[i] = 0;
+        }
+
         // accelerate edges that connect simple bodies with no rotations
-        if (m->body_simple[b1] == 2 && m->body_simple[b2] == 2) {
+        else if (m->body_simple[b1] == 2 && m->body_simple[b2] == 2) {
           m->flexedge_invweight0[i] = (1/m->body_mass[b1] + 1/m->body_mass[b2])/2;
         }
 

src/user/user_model.cc

@@ -1569,8 +1569,8 @@ void mjCModel::CopyTree(mjModel* m) {
       qposadr += nPOS[pj->type];
     }
 
-    // simple body with no rotational dofs: promote to simple level 2
+    // simple body with sliders and no rotational dofs: promote to simple level 2
-    if (m->body_simple[i]) {
+    if (m->body_simple[i] && m->body_dofnum[i]) {
       m->body_simple[i] = 2;
       for (int j=0; j<(int)pb->joints.size(); j++) {
         if (pb->joints[j]->type!=mjJNT_SLIDE) {
@@ -1979,10 +1979,15 @@ void mjCModel::CopyObjects(mjModel* m) {
       memcpy(m->flex_vertbodyid + vert_adr, pfl->vertbodyid.data(), pfl->nvert*sizeof(int));
     }
 
-    // convert edge pairs to int array
+    // convert edge pairs to int array, set edge rigid
-    for (int i=0; i<pfl->nedge; i++) {
+    for (int k=0; k<pfl->nedge; k++) {
-      m->flex_edge[2*(edge_adr+i)] = pfl->edge[i].first;
+      m->flex_edge[2*(edge_adr+k)] = pfl->edge[k].first;
-      m->flex_edge[2*(edge_adr+i)+1] = pfl->edge[i].second;
+      m->flex_edge[2*(edge_adr+k)+1] = pfl->edge[k].second;
+
+      // check if vertex body weldids are the same
+      int b1 = pfl->vertbodyid[pfl->edge[k].first];
+      int b2 = pfl->vertbodyid[pfl->edge[k].second];
+      m->flexedge_rigid[edge_adr+k] = (bodies[b1]->weldid == bodies[b2]->weldid);
     }
 
     // advance counters

unity/Runtime/Bindings/MjBindings.cs

@@ -5337,6 +5337,7 @@ public unsafe struct mjModel_ {
   public double* flex_edgedamping;
   public byte* flex_edgeequality;
   public byte* flex_rigid;
+  public byte* flexedge_rigid;
   public byte* flex_centered;
   public byte* flex_flatskin;
   public int* flex_bvhadr;