Merge pull request #236 from thowell/tendon_friction

add tendon friction

Author: thowell

mujoco_warp/_src/constraint.py

@@ -437,7 +437,7 @@ def _efc_equality_tendon(
 
 
 @wp.kernel
-def _efc_friction(
+def _efc_friction_dof(
   # Model:
   opt_timestep: float,
   dof_invweight0: wp.array2d(dtype=float),
@@ -461,7 +461,6 @@ def _efc_friction(
   efc_aref_out: wp.array(dtype=float),
   efc_frictionloss_out: wp.array(dtype=float),
 ):
-  # TODO(team): tendon
   worldid, dofid = wp.tid()
 
   if dof_frictionloss[worldid, dofid] <= 0.0:
@@ -500,6 +499,72 @@ def _efc_friction(
   )
 
 
+@wp.kernel
+def _efc_friction_tendon(
+  # Model:
+  nv: int,
+  opt_timestep: float,
+  tendon_solref_fri: wp.array2d(dtype=wp.vec2),
+  tendon_solimp_fri: wp.array2d(dtype=vec5),
+  tendon_frictionloss: wp.array2d(dtype=float),
+  tendon_invweight0: wp.array2d(dtype=float),
+  # Data in:
+  qvel_in: wp.array2d(dtype=float),
+  ten_J_in: wp.array3d(dtype=float),
+  # In:
+  refsafe_in: int,
+  # Data out:
+  nf_out: wp.array(dtype=int),
+  nefc_out: wp.array(dtype=int),
+  efc_worldid_out: wp.array(dtype=int),
+  efc_id_out: wp.array(dtype=int),
+  efc_J_out: wp.array2d(dtype=float),
+  efc_pos_out: wp.array(dtype=float),
+  efc_margin_out: wp.array(dtype=float),
+  efc_D_out: wp.array(dtype=float),
+  efc_aref_out: wp.array(dtype=float),
+  efc_frictionloss_out: wp.array(dtype=float),
+):
+  worldid, tenid = wp.tid()
+
+  frictionloss = tendon_frictionloss[worldid, tenid]
+  if frictionloss <= 0.0:
+    return
+
+  efcid = wp.atomic_add(nefc_out, 0, 1)
+  wp.atomic_add(nf_out, 0, 1)
+  efc_worldid_out[efcid] = worldid
+
+  Jqvel = float(0.0)
+
+  # TODO(team): parallelize
+  for i in range(nv):
+    J = ten_J_in[worldid, tenid, i]
+    efc_J_out[efcid, i] = J
+    Jqvel += J * qvel_in[worldid, i]
+
+  _update_efc_row(
+    opt_timestep,
+    refsafe_in,
+    efcid,
+    0.0,
+    0.0,
+    tendon_invweight0[worldid, tenid],
+    tendon_solref_fri[worldid, tenid],
+    tendon_solimp_fri[worldid, tenid],
+    0.0,
+    Jqvel,
+    frictionloss,
+    tenid,
+    efc_id_out,
+    efc_pos_out,
+    efc_margin_out,
+    efc_D_out,
+    efc_aref_out,
+    efc_frictionloss_out,
+  )
+
+
 @wp.kernel
 def _efc_equality_weld(
   # Model:
@@ -1477,7 +1542,7 @@ def make_constraint(m: types.Model, d: types.Data):
 
     if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS.value):
       wp.launch(
-        _efc_friction,
+        _efc_friction_dof,
         dim=(d.nworld, m.nv),
         inputs=[
           m.opt.timestep,
@@ -1503,6 +1568,34 @@ def make_constraint(m: types.Model, d: types.Data):
         ],
       )
 
+      wp.launch(
+        _efc_friction_tendon,
+        dim=(d.nworld, m.ntendon),
+        inputs=[
+          m.nv,
+          m.opt.timestep,
+          m.tendon_solref_fri,
+          m.tendon_solimp_fri,
+          m.tendon_frictionloss,
+          m.tendon_invweight0,
+          d.qvel,
+          d.ten_J,
+          refsafe,
+        ],
+        outputs=[
+          d.nf,
+          d.nefc,
+          d.efc.worldid,
+          d.efc.id,
+          d.efc.J,
+          d.efc.pos,
+          d.efc.margin,
+          d.efc.D,
+          d.efc.aref,
+          d.efc.frictionloss,
+        ],
+      )
+
     # limit
     if not (m.opt.disableflags & types.DisableBit.LIMIT.value):
       limit_ball = m.jnt_limited_ball_adr.size > 0

mujoco_warp/_src/io.py

@@ -44,9 +44,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   if mjm.nflex > 1:
     raise NotImplementedError("Only one flex is unsupported.")
 
-  if mjm.tendon_frictionloss.any():
-    raise NotImplementedError("Tendon frictionloss is unsupported.")
-
   if mjm.geom_fluid.any():
     raise NotImplementedError("Ellipsoid fluid model not implemented.")
 
@@ -503,8 +500,11 @@ def create_nmodel_batched_array(mjm_array, dtype):
     tendon_limited_adr=wp.array(np.nonzero(mjm.tendon_limited)[0], dtype=wp.int32, ndim=1),
     tendon_solref_lim=create_nmodel_batched_array(mjm.tendon_solref_lim, dtype=wp.vec2f),
     tendon_solimp_lim=create_nmodel_batched_array(mjm.tendon_solimp_lim, dtype=types.vec5),
+    tendon_solref_fri=create_nmodel_batched_array(mjm.tendon_solref_fri, dtype=wp.vec2f),
+    tendon_solimp_fri=create_nmodel_batched_array(mjm.tendon_solimp_fri, dtype=types.vec5),
     tendon_range=create_nmodel_batched_array(mjm.tendon_range, dtype=wp.vec2f),
     tendon_margin=create_nmodel_batched_array(mjm.tendon_margin, dtype=float),
+    tendon_frictionloss=create_nmodel_batched_array(mjm.tendon_frictionloss, dtype=float),
     tendon_length0=create_nmodel_batched_array(mjm.tendon_length0, dtype=float),
     tendon_invweight0=create_nmodel_batched_array(mjm.tendon_invweight0, dtype=float),
     wrap_objid=wp.array(mjm.wrap_objid, dtype=int),

mujoco_warp/_src/types.py

@@ -769,8 +769,11 @@ class Model:
     tendon_limited_adr: addresses for limited tendons        (<=ntendon,)
     tendon_solref_lim: constraint solver reference: limit    (nworld, ntendon, mjNREF)
     tendon_solimp_lim: constraint solver impedance: limit    (nworld, ntendon, mjNIMP)
+    tendon_solref_fri: constraint solver reference: friction (nworld, ntendon, mjNREF)
+    tendon_solimp_fri: constraint solver impedance: friction (nworld, ntendon, mjNIMP)
     tendon_range: tendon length limits                       (nworld, ntendon, 2)
     tendon_margin: min distance for limit detection          (nworld, ntendon,)
+    tendon_frictionloss: loss due to friction                (nworld, ntendon,)
     tendon_length0: tendon length in qpos0                   (nworld, ntendon,)
     tendon_invweight0: inv. weight in qpos0                  (nworld, ntendon,)
     wrap_objid: object id: geom, site, joint                 (nwrap,)
@@ -1011,10 +1014,14 @@ class Model:
   tendon_num: wp.array(dtype=int)
   tendon_limited: wp.array(dtype=int)
   tendon_limited_adr: wp.array(dtype=int)
+
   tendon_solref_lim: wp.array2d(dtype=wp.vec2)
   tendon_solimp_lim: wp.array2d(dtype=vec5)
+  tendon_solref_fri: wp.array2d(dtype=wp.vec2)
+  tendon_solimp_fri: wp.array2d(dtype=vec5)
   tendon_range: wp.array2d(dtype=wp.vec2)
   tendon_margin: wp.array2d(dtype=float)
+  tendon_frictionloss: wp.array2d(dtype=float)
   tendon_length0: wp.array2d(dtype=float)
   tendon_invweight0: wp.array2d(dtype=float)
   wrap_objid: wp.array(dtype=int)

mujoco_warp/test_data/constraints.xml

@@ -77,6 +77,13 @@
     </body>
   </worldbody>
 
+  <tendon>
+    <fixed name="tendon_1" frictionloss=".1">
+      <joint joint="joint3" coef=".1"/>
+      <joint joint="joint4" coef="-.2"/>
+    </fixed>
+  </tendon>
+
   <equality>
     <connect name="c_site"   site1="site0"      site2="site1"/>
     <connect name="connect"  body1="anchor1"    body2="beam1"   anchor="1 0 -1" />