Update flex_bending dimensions (#628)

* update flex_bending dimensions

* add missing flatten

* make format uglier thanks to ruff

* add curved shell support

Author: quagla

mujoco_warp/_src/io.py

@@ -610,7 +610,7 @@ def create_nmodel_batched_array(mjm_array, dtype, expand_dim=True):
     flex_elemedge=wp.array(mjm.flex_elemedge, dtype=int),
     flexedge_length0=wp.array(mjm.flexedge_length0, dtype=float),
     flex_stiffness=wp.array(mjm.flex_stiffness.flatten(), dtype=float),
-    flex_bending=wp.array(mjm.flex_bending, dtype=wp.mat44f),
+    flex_bending=wp.array(mjm.flex_bending.flatten(), dtype=float),
     flex_damping=wp.array(mjm.flex_damping, dtype=float),
     mesh_vertadr=wp.array(mjm.mesh_vertadr, dtype=int),
     mesh_vertnum=wp.array(mjm.mesh_vertnum, dtype=int),

mujoco_warp/_src/passive.py

@@ -454,7 +454,7 @@ def _flex_bending(
   flex_vertbodyid: wp.array(dtype=int),
   flex_edge: wp.array(dtype=wp.vec2i),
   flex_edgeflap: wp.array(dtype=wp.vec2i),
-  flex_bending: wp.array(dtype=wp.mat44f),
+  flex_bending: wp.array(dtype=float),
   # Data in:
   flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
   # Data out:
@@ -477,11 +477,23 @@ def _flex_bending(
   if v[3] == -1:
     return
 
+  frc = wp.mat(0.0, shape=(4, 3))
+  if flex_bending[17 * edgeid + 16]:
+    v0 = flexvert_xpos_in[worldid, v[0]]
+    v1 = flexvert_xpos_in[worldid, v[1]]
+    v2 = flexvert_xpos_in[worldid, v[2]]
+    v3 = flexvert_xpos_in[worldid, v[3]]
+    frc[1, :] = wp.cross(v2 - v0, v3 - v0)
+    frc[2, :] = wp.cross(v3 - v0, v1 - v0)
+    frc[3, :] = wp.cross(v1 - v0, v2 - v0)
+    frc[0, :] = -(frc[1, :] + frc[2, :] + frc[3, :])
+
   force = wp.mat(0.0, shape=(nvert, 3))
   for i in range(nvert):
-    for j in range(nvert):
-      for x in range(3):
-        force[i, x] -= flex_bending[edgeid][i, j] * flexvert_xpos_in[worldid, v[j]][x]
+    for x in range(3):
+      for j in range(nvert):
+        force[i, x] -= flex_bending[17 * edgeid + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
+    force[i, x] -= flex_bending[17 * edgeid + 16] * frc[i, x]
 
   for i in range(nvert):
     bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]]

mujoco_warp/_src/types.py

@@ -1201,7 +1201,7 @@ class Model:
   flex_elemedge: wp.array(dtype=int)
   flexedge_length0: wp.array(dtype=float)
   flex_stiffness: wp.array(dtype=float)
-  flex_bending: wp.array(dtype=wp.mat44f)
+  flex_bending: wp.array(dtype=float)
   flex_damping: wp.array(dtype=float)
   mesh_vertadr: wp.array(dtype=int)
   mesh_vertnum: wp.array(dtype=int)