add tendon spring damper

mujoco_warp/_src/io.py

@@ -469,6 +469,9 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     tendon_solimp_lim=wp.array(mjm.tendon_solimp_lim, dtype=types.vec5),
     tendon_range=wp.array(mjm.tendon_range, dtype=wp.vec2f),
     tendon_margin=wp.array(mjm.tendon_margin, dtype=float),
+    tendon_stiffness=wp.array(mjm.tendon_stiffness, dtype=float),
+    tendon_damping=wp.array(mjm.tendon_damping, dtype=float),
+    tendon_lengthspring=wp.array(mjm.tendon_lengthspring, dtype=wp.vec2),
     tendon_length0=wp.array(mjm.tendon_length0, dtype=float),
     tendon_invweight0=wp.array(mjm.tendon_invweight0, dtype=float),
     wrap_objid=wp.array(mjm.wrap_objid, dtype=int),

mujoco_warp/_src/passive.py

@@ -25,29 +25,34 @@
 
 
 @wp.kernel
-def _spring_passive(
+def _spring_damper_dof_passive(
   # Model:
   qpos_spring: wp.array(dtype=float),
   jnt_type: wp.array(dtype=int),
   jnt_qposadr: wp.array(dtype=int),
   jnt_dofadr: wp.array(dtype=int),
   jnt_stiffness: wp.array(dtype=float),
+  dof_damping: wp.array(dtype=float),
   # Data in:
   qpos_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d(dtype=float),
   # Data out:
   qfrc_spring_out: wp.array2d(dtype=float),
+  qfrc_damper_out: wp.array2d(dtype=float),
 ):
   worldid, jntid = wp.tid()
   stiffness = jnt_stiffness[jntid]
   dofid = jnt_dofadr[jntid]
+  damping = dof_damping[dofid]
 
-  if stiffness == 0.0:
+  if stiffness == 0.0 and damping == 0.0:
     return
 
   jnttype = jnt_type[jntid]
   qposid = jnt_qposadr[jntid]
 
   if jnttype == wp.static(JointType.FREE.value):
+    # spring
     dif = wp.vec3(
       qpos_in[worldid, qposid + 0] - qpos_spring[qposid + 0],
       qpos_in[worldid, qposid + 1] - qpos_spring[qposid + 1],
@@ -72,7 +77,16 @@ def _spring_passive(
     qfrc_spring_out[worldid, dofid + 3] = -stiffness * dif[0]
     qfrc_spring_out[worldid, dofid + 4] = -stiffness * dif[1]
     qfrc_spring_out[worldid, dofid + 5] = -stiffness * dif[2]
+
+    # damper
+    qfrc_damper_out[worldid, dofid + 0] = -damping * qvel_in[worldid, dofid + 0]
+    qfrc_damper_out[worldid, dofid + 1] = -damping * qvel_in[worldid, dofid + 1]
+    qfrc_damper_out[worldid, dofid + 2] = -damping * qvel_in[worldid, dofid + 2]
+    qfrc_damper_out[worldid, dofid + 3] = -damping * qvel_in[worldid, dofid + 3]
+    qfrc_damper_out[worldid, dofid + 4] = -damping * qvel_in[worldid, dofid + 4]
+    qfrc_damper_out[worldid, dofid + 5] = -damping * qvel_in[worldid, dofid + 5]
   elif jnttype == wp.static(JointType.BALL.value):
+    # spring
     rot = wp.quat(
       qpos_in[worldid, qposid + 0],
       qpos_in[worldid, qposid + 1],
@@ -89,28 +103,67 @@ def _spring_passive(
     qfrc_spring_out[worldid, dofid + 0] = -stiffness * dif[0]
     qfrc_spring_out[worldid, dofid + 1] = -stiffness * dif[1]
     qfrc_spring_out[worldid, dofid + 2] = -stiffness * dif[2]
+
+    # damper
+    qfrc_damper_out[worldid, dofid + 0] = -damping * qvel_in[worldid, dofid + 0]
+    qfrc_damper_out[worldid, dofid + 1] = -damping * qvel_in[worldid, dofid + 1]
+    qfrc_damper_out[worldid, dofid + 2] = -damping * qvel_in[worldid, dofid + 2]
   else:  # mjJNT_SLIDE, mjJNT_HINGE
+    # spring
     fdif = qpos_in[worldid, qposid] - qpos_spring[qposid]
     qfrc_spring_out[worldid, dofid] = -stiffness * fdif
 
+    # damper
+    qfrc_damper_out[worldid, dofid] = -damping * qvel_in[worldid, dofid]
+
 
 @wp.kernel
-def _damper_passive(
+def _spring_damper_tendon_passive(
   # Model:
-  dof_damping: wp.array(dtype=float),
+  tendon_stiffness: wp.array(dtype=float),
+  tendon_damping: wp.array(dtype=float),
+  tendon_lengthspring: wp.array(dtype=wp.vec2),
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  qfrc_spring_in: wp.array2d(dtype=float),
+  ten_velocity_in: wp.array2d(dtype=float),
+  ten_length_in: wp.array2d(dtype=float),
+  ten_J_in: wp.array3d(dtype=float),
   # Data out:
+  qfrc_spring_out: wp.array2d(dtype=float),
   qfrc_damper_out: wp.array2d(dtype=float),
-  qfrc_passive_out: wp.array2d(dtype=float),
 ):
-  worldid, dofid = wp.tid()
+  worldid, tenid, dofid = wp.tid()
+
+  stiffness = tendon_stiffness[tenid]
+  damping = tendon_damping[tenid]
+
+  if stiffness == 0.0 and damping == 0.0:
+    return
+
+  J = ten_J_in[worldid, tenid, dofid]
 
-  qfrc_damper = -dof_damping[dofid] * qvel_in[worldid, dofid]
+  if stiffness:
+    # compute spring force along tendon
+    length = ten_length_in[worldid, tenid]
+    lengthspring = tendon_lengthspring[tenid]
+    lower = lengthspring[0]
+    upper = lengthspring[1]
 
-  qfrc_damper_out[worldid, dofid] = qfrc_damper
-  qfrc_passive_out[worldid, dofid] = qfrc_damper + qfrc_spring_in[worldid, dofid]
+    if length > upper:
+      frc_spring = stiffness * (upper - length)
+    elif length < lower:
+      frc_spring = stiffness * (lower - length)
+    else:
+      frc_spring = 0.0
+
+    # transform to joint torque
+    wp.atomic_add(qfrc_spring_out[worldid], dofid, J * frc_spring)
+
+  if damping:
+    # compute damper linear force along tendon
+    frc_damper = -damping * ten_velocity_in[worldid, tenid]
+
+    # transform to joint torque
+    wp.atomic_add(qfrc_damper_out[worldid], dofid, J * frc_damper)
 
 
 @wp.kernel
@@ -143,50 +196,82 @@ def _gravity_force(
 
 
 @wp.kernel
-def _qfrc_passive_gravcomp(
+def _qfrc_passive(
   # Model:
   jnt_actgravcomp: wp.array(dtype=int),
   dof_jntid: wp.array(dtype=int),
   # Data in:
+  qfrc_spring_in: wp.array2d(dtype=float),
+  qfrc_damper_in: wp.array2d(dtype=float),
   qfrc_gravcomp_in: wp.array2d(dtype=float),
+  # In:
+  gravcomp: bool,
   # Data out:
   qfrc_passive_out: wp.array2d(dtype=float),
 ):
   worldid, dofid = wp.tid()
+  qfrc_passive = qfrc_spring_in[worldid, dofid]
+  qfrc_passive += qfrc_damper_in[worldid, dofid]
 
   # add gravcomp unless added by actuators
-  if jnt_actgravcomp[dof_jntid[dofid]]:
-    return
+  if gravcomp and not jnt_actgravcomp[dof_jntid[dofid]]:
+    qfrc_passive += qfrc_gravcomp_in[worldid, dofid]
 
-  qfrc_passive_out[worldid, dofid] += qfrc_gravcomp_in[worldid, dofid]
+  qfrc_passive_out[worldid, dofid] = qfrc_passive
 
 
 @event_scope
 def passive(m: Model, d: Data):
   """Adds all passive forces."""
   if m.opt.disableflags & DisableBit.PASSIVE:
-    d.qfrc_passive.zero_()
+    d.qfrc_spring.zero_()
+    d.qfrc_damper.zero_()
     d.qfrc_gravcomp.zero_()
+    d.qfrc_passive.zero_()
     return
 
   # TODO(team): mj_ellipsoidFluidModel
   # TODO(team): mj_inertiaBoxFluidModell
 
   d.qfrc_spring.zero_()
+  d.qfrc_damper.zero_()
   wp.launch(
-    _spring_passive,
+    _spring_damper_dof_passive,
     dim=(d.nworld, m.njnt),
-    inputs=[m.qpos_spring, m.jnt_type, m.jnt_qposadr, m.jnt_dofadr, m.jnt_stiffness, d.qpos],
-    outputs=[d.qfrc_spring],
-  )
-  wp.launch(
-    _damper_passive,
-    dim=(d.nworld, m.nv),
-    inputs=[m.dof_damping, d.qvel, d.qfrc_spring],
-    outputs=[d.qfrc_damper, d.qfrc_passive],
+    inputs=[
+      m.qpos_spring,
+      m.jnt_type,
+      m.jnt_qposadr,
+      m.jnt_dofadr,
+      m.jnt_stiffness,
+      m.dof_damping,
+      d.qpos,
+      d.qvel,
+    ],
+    outputs=[d.qfrc_spring, d.qfrc_damper],
   )
-  if m.ngravcomp and not (m.opt.disableflags & DisableBit.GRAVITY):
-    d.qfrc_gravcomp.zero_()
+
+  if m.ntendon:
+    wp.launch(
+      _spring_damper_tendon_passive,
+      dim=(d.nworld, m.ntendon, m.nv),
+      inputs=[
+        m.tendon_stiffness,
+        m.tendon_damping,
+        m.tendon_lengthspring,
+        d.ten_velocity,
+        d.ten_length,
+        d.ten_J,
+      ],
+      outputs=[
+        d.qfrc_spring,
+        d.qfrc_damper,
+      ],
+    )
+
+  d.qfrc_gravcomp.zero_()
+  gravcomp = m.ngravcomp and not (m.opt.disableflags & DisableBit.GRAVITY)
+  if gravcomp:
     wp.launch(
       _gravity_force,
       dim=(d.nworld, m.nbody - 1, m.nv),
@@ -203,9 +288,10 @@ def passive(m: Model, d: Data):
       ],
       outputs=[d.qfrc_gravcomp],
     )
-    wp.launch(
-      _qfrc_passive_gravcomp,
-      dim=(d.nworld, m.nv),
-      inputs=[m.jnt_actgravcomp, m.dof_jntid, d.qfrc_gravcomp],
-      outputs=[d.qfrc_passive],
-    )
+
+  wp.launch(
+    _qfrc_passive,
+    dim=(d.nworld, m.nv),
+    inputs=[m.jnt_actgravcomp, m.dof_jntid, d.qfrc_spring, d.qfrc_damper, d.qfrc_gravcomp, gravcomp],
+    outputs=[d.qfrc_passive],
+  )

mujoco_warp/_src/passive_test.py

@@ -41,7 +41,7 @@ def test_passive(self, gravity):
     """Tests passive."""
     _, mjd, m, d = test_util.fixture("pendula.xml", gravity=gravity)
 
-    for arr in (d.qfrc_spring, d.qfrc_damper, d.qfrc_passive):
+    for arr in (d.qfrc_spring, d.qfrc_damper, d.qfrc_gravcomp, d.qfrc_passive):
       arr.zero_()
 
     mjwarp.passive(m, d)

mujoco_warp/_src/types.py

@@ -748,6 +748,9 @@ class Model:
     tendon_solimp_lim: constraint solver impedance: limit    (ntendon, mjNIMP)
     tendon_range: tendon length limits                       (ntendon, 2)
     tendon_margin: min distance for limit detection          (ntendon,)
+    tendon_stiffness: stiffness coefficient                  (ntendon,)
+    tendon_damping: damping coefficient                      (ntendon,)
+    tendon_lengthspring: spring resting length range         (ntendon, 2)
     tendon_length0: tendon length in qpos0                   (ntendon,)
     tendon_invweight0: inv. weight in qpos0                  (ntendon,)
     wrap_objid: object id: geom, site, joint                 (nwrap,)
@@ -961,6 +964,9 @@ class Model:
   tendon_solimp_lim: wp.array(dtype=vec5)
   tendon_range: wp.array(dtype=wp.vec2)
   tendon_margin: wp.array(dtype=float)
+  tendon_stiffness: wp.array(dtype=float)
+  tendon_damping: wp.array(dtype=float)
+  tendon_lengthspring: wp.array(dtype=wp.vec2)
   tendon_length0: wp.array(dtype=float)
   tendon_invweight0: wp.array(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" stiffness=".1" damping=".2">
+      <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" />