Bending stiffness for curved shells in flex.

PiperOrigin-RevId: 796361841
Change-Id: I105ea235fe2a8e1bcbde67276d77fa92eed214a1

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -609,7 +609,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     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),

mjx/mujoco/mjx/third_party/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:
@@ -481,7 +481,7 @@ def _flex_bending(
   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]
+        force[i, x] -= flex_bending[17*edgeid + 4*i + j] * flexvert_xpos_in[worldid, v[j]][x]
 
   for i in range(nvert):
     bodyid = flex_vertbodyid[flex_vertadr[f] + v[i]]

mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py

@@ -1199,7 +1199,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)

mjx/mujoco/mjx/warp/forward.py

@@ -124,7 +124,7 @@ def _forward_shim(
     eq_solref: wp.array2d(dtype=wp.vec2),
     eq_ten_adr: wp.array(dtype=int),
     eq_wld_adr: wp.array(dtype=int),
-    flex_bending: wp.array(dtype=wp.mat44f),
+    flex_bending: wp.array(dtype=float),
     flex_damping: wp.array(dtype=float),
     flex_dim: wp.array(dtype=int),
     flex_edge: wp.array(dtype=wp.vec2i),
@@ -2148,7 +2148,7 @@ def _step_shim(
     eq_solref: wp.array2d(dtype=wp.vec2),
     eq_ten_adr: wp.array(dtype=int),
     eq_wld_adr: wp.array(dtype=int),
-    flex_bending: wp.array(dtype=wp.mat44f),
+    flex_bending: wp.array(dtype=float),
     flex_damping: wp.array(dtype=float),
     flex_dim: wp.array(dtype=int),
     flex_edge: wp.array(dtype=wp.vec2i),

mjx/mujoco/mjx/warp/types.py

@@ -737,7 +737,7 @@ _NDIM = {
         'eq_type': 1,
         'eq_wld_adr': 1,
         'exclude_signature': 1,
-        'flex_bending': 3,
+        'flex_bending': 1,
         'flex_damping': 1,
         'flex_dim': 1,
         'flex_edge': 2,