update _flex_edges (google-deepmind#1208)

Author: thowell

mujoco_warp/_src/smooth.py

@@ -242,10 +242,8 @@ def _flex_vertices(
 def _flex_edges(
   # Model:
   nflex: int,
-  body_parentid: wp.array(dtype=int),
   body_rootid: wp.array(dtype=int),
   body_dofadr: wp.array(dtype=int),
-  dof_bodyid: wp.array(dtype=int),
   flex_vertadr: wp.array(dtype=int),
   flex_edgeadr: wp.array(dtype=int),
   flex_edgenum: wp.array(dtype=int),
@@ -286,64 +284,32 @@ def _flex_edges(
 
   dofi = body_dofadr[b1]
   dofj = body_dofadr[b2]
-  dofi0 = dofi + 0
-  dofi1 = dofi + 1
-  dofi2 = dofi + 2
-  dofj0 = dofj + 0
-  dofj1 = dofj + 1
-  dofj2 = dofj + 2
-
-  vel1 = wp.vec3(qvel_in[worldid, dofi0], qvel_in[worldid, dofi1], qvel_in[worldid, dofi2])
-  vel2 = wp.vec3(qvel_in[worldid, dofj0], qvel_in[worldid, dofj1], qvel_in[worldid, dofj2])
+
+  vel1 = wp.vec3(qvel_in[worldid, dofi], qvel_in[worldid, dofi + 1], qvel_in[worldid, dofi + 2])
+  vel2 = wp.vec3(qvel_in[worldid, dofj], qvel_in[worldid, dofj + 1], qvel_in[worldid, dofj + 2])
   flexedge_velocity_out[worldid, edgeid] = wp.dot(vel2 - vel1, edge)
 
   rowadr = flexedge_J_rowadr[edgeid]
 
-  sparseid0 = rowadr + 0
-  sparseid1 = rowadr + 1
-  sparseid2 = rowadr + 2
-  sparseid3 = rowadr + 3
-  sparseid4 = rowadr + 4
-  sparseid5 = rowadr + 5
-
-  # TODO(team): jacdif
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi0, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi0, worldid)
-  jacdif = jacp2 - jacp1
-  Ji0 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi1, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi1, worldid)
-  jacdif = jacp2 - jacp1
-  Ji1 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofi2, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofi2, worldid)
-  jacdif = jacp2 - jacp1
-  Ji2 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj0, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj0, worldid)
-  jacdif = jacp2 - jacp1
-  Jj0 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj1, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj1, worldid)
-  jacdif = jacp2 - jacp1
-  Jj1 = wp.dot(jacdif, edge)
-
-  jacp1, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos1, b1, dofj2, worldid)
-  jacp2, _ = support.jac_dof(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos2, b2, dofj2, worldid)
-  jacdif = jacp2 - jacp1
-  Jj2 = wp.dot(jacdif, edge)
-
-  flexedge_J_out[worldid, sparseid0] = Ji0
-  flexedge_J_out[worldid, sparseid1] = Ji1
-  flexedge_J_out[worldid, sparseid2] = Ji2
-  flexedge_J_out[worldid, sparseid3] = Jj0
-  flexedge_J_out[worldid, sparseid4] = Jj1
-  flexedge_J_out[worldid, sparseid5] = Jj2
+  # compute offsets once per body (avoids 12 redundant tree-ancestry walks in jac_dof)
+  offset1 = pos1 - wp.vec3(subtree_com_in[worldid, body_rootid[b1]])
+  offset2 = pos2 - wp.vec3(subtree_com_in[worldid, body_rootid[b2]])
+
+  # body1 DOFs: b1 is in subtree, b2 is not -> jacdif = 0 - jacp1 = -jacp1
+  for k in range(3):
+    cdof = cdof_in[worldid, dofi + k]
+    cdof_ang = wp.spatial_top(cdof)
+    cdof_lin = wp.spatial_bottom(cdof)
+    jacp1 = cdof_lin + wp.cross(cdof_ang, offset1)
+    flexedge_J_out[worldid, rowadr + k] = wp.dot(-jacp1, edge)
+
+  # body2 DOFs: b2 is in subtree, b1 is not -> jacdif = jacp2 - 0 = jacp2
+  for k in range(3):
+    cdof = cdof_in[worldid, dofj + k]
+    cdof_ang = wp.spatial_top(cdof)
+    cdof_lin = wp.spatial_bottom(cdof)
+    jacp2 = cdof_lin + wp.cross(cdof_ang, offset2)
+    flexedge_J_out[worldid, rowadr + 3 + k] = wp.dot(jacp2, edge)
 
 
 @event_scope
@@ -415,10 +381,8 @@ def flex(m: Model, d: Data):
     dim=(d.nworld, m.nflexedge),
     inputs=[
       m.nflex,
-      m.body_parentid,
       m.body_rootid,
       m.body_dofadr,
-      m.dof_bodyid,
       m.flex_vertadr,
       m.flex_edgeadr,
       m.flex_edgenum,