Fix multi flex indexing

Author: StafaH

mujoco_warp/_src/collision_flex.py

@@ -397,6 +397,7 @@ def _flex_narrowphase_dim2(
   flex_dim: wp.array(dtype=int),
   flex_vertadr: wp.array(dtype=int),
   flex_elemadr: wp.array(dtype=int),
+  flex_elemdataadr: wp.array(dtype=int),
   flex_elemnum: wp.array(dtype=int),
   flex_elem: wp.array(dtype=int),
   flex_radius: wp.array(dtype=float),
@@ -443,7 +444,7 @@ def _flex_narrowphase_dim2(
   tri_radius = flex_radius[flexid]
   tri_margin = flex_margin[flexid]
 
-  elem_data_idx = elemid * 3
+  elem_data_idx = flex_elemdataadr[flexid] + (elemid - flex_elemadr[flexid]) * 3
   v0_local = flex_elem[elem_data_idx]
   v1_local = flex_elem[elem_data_idx + 1]
   v2_local = flex_elem[elem_data_idx + 2]
@@ -708,6 +709,7 @@ def flex_narrowphase(m: Model, d: Data):
       m.flex_dim,
       m.flex_vertadr,
       m.flex_elemadr,
+      m.flex_elemdataadr,
       m.flex_elemnum,
       m.flex_elem,
       m.flex_radius,

mujoco_warp/_src/passive.py

@@ -573,6 +573,7 @@ def _flex_elasticity(
   flex_vertadr: wp.array(dtype=int),
   flex_edgeadr: wp.array(dtype=int),
   flex_elemadr: wp.array(dtype=int),
+  flex_elemdataadr: wp.array(dtype=int),
   flex_elemnum: wp.array(dtype=int),
   flex_elemedgeadr: wp.array(dtype=int),
   flex_vertbodyid: wp.array(dtype=int),
@@ -593,12 +594,14 @@ def _flex_elasticity(
   worldid, elemid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
+  f = int(0)
   for i in range(nflex):
     locid = elemid - flex_elemadr[i]
     if locid >= 0 and locid < flex_elemnum[i]:
       f = i
       break
 
+  local_elemid = elemid - flex_elemadr[f]
   dim = flex_dim[f]
   nvert = dim + 1
   nedge = nvert * (nvert - 1) / 2
@@ -612,10 +615,11 @@ def _flex_elasticity(
   else:
     kD = 0.0
 
+  elem_data_adr = flex_elemdataadr[f] + local_elemid * (dim + 1)
   gradient = wp.matrix(0.0, shape=(6, 6))
   for e in range(nedge):
-    vert0 = flex_elem[(dim + 1) * elemid + edges[e, 0]]
-    vert1 = flex_elem[(dim + 1) * elemid + edges[e, 1]]
+    vert0 = flex_elem[elem_data_adr + edges[e, 0]]
+    vert1 = flex_elem[elem_data_adr + edges[e, 1]]
     xpos0 = flexvert_xpos_in[worldid, vert0]
     xpos1 = flexvert_xpos_in[worldid, vert1]
     for i in range(3):
@@ -624,7 +628,7 @@ def _flex_elasticity(
 
   elongation = wp.spatial_vectorf(0.0)
   for e in range(nedge):
-    idx = flex_elemedge[elemid * nedge + e]
+    idx = flex_elemedge[flex_elemedgeadr[f] + local_elemid * int(nedge) + e]
     vel = flexedge_velocity_in[worldid, flex_edgeadr[f] + idx]
     deformed = flexedge_length_in[worldid, flex_edgeadr[f] + idx]
     reference = flexedge_length0[flex_edgeadr[f] + idx]
@@ -647,7 +651,7 @@ def _flex_elasticity(
           force[edges[ed2, i], x] -= elongation[ed1] * gradient[ed2, 3 * i + x] * metric[ed1, ed2]
 
   for v in range(nvert):
-    vert = flex_elem[(dim + 1) * elemid + v]
+    vert = flex_elem[elem_data_adr + v]
     bodyid = flex_vertbodyid[flex_vertadr[f] + vert]
     for x in range(3):
       wp.atomic_add(qfrc_spring_out, worldid, body_dofadr[bodyid] + x, force[v, x])
@@ -783,6 +787,7 @@ def passive(m: Model, d: Data):
         m.flex_vertadr,
         m.flex_edgeadr,
         m.flex_elemadr,
+        m.flex_elemdataadr,
         m.flex_elemnum,
         m.flex_elemedgeadr,
         m.flex_vertbodyid,

mujoco_warp/_src/types.py

@@ -1365,6 +1365,7 @@ class Model:
   flex_edgeadr: array("nflex", int)
   flex_edgenum: array("nflex", int)
   flex_elemadr: array("nflex", int)
+  flex_elemdataadr: array("nflex", int)
   flex_elemnum: array("nflex", int)
   flex_elemedgeadr: array("nflex", int)
   flex_shellnum: array("nflex", int)