Add support for joint and tendon frictionloss to MJX.

PiperOrigin-RevId: 670679777
Change-Id: I2eb56d013c11d4e1851f1c2edd6d98faaf854db9

doc/changelog.rst

@@ -45,6 +45,7 @@ MJX
 - Added support for :ref:`implicitfast integration<geIntegration>` for all cases except
   :doc:`fluid drag <computation/fluid>`.
 - Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices.
+- Added support for joint and tendon :ref:`frictionloss <coFriction>`.
 
 Bug fixes
 ^^^^^^^^^

doc/mjx.rst

@@ -199,7 +199,7 @@ The following features are **fully supported** in MJX:
      - ``PLANE``, ``HFIELD``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` are fully implemented. ``ELLIPSOID`` and
        ``CYLINDER`` are implemented but only collide with other primitives, note that ``BOX`` is implemented as a mesh.
    * - :ref:`Constraint <mjtConstraint>`
-     - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL``, ``CONTACT_ELLIPTIC``
+     - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``, ``FRICTION_TENDON``
    * - :ref:`Equality <mjtEq>`
      - ``CONNECT``, ``WELD``, ``JOINT``, ``TENDON``
    * - :ref:`Integrator <mjtIntegrator>`
@@ -228,8 +228,6 @@ The following features are **in development** and coming soon:
    * - :ref:`Geom <mjtGeom>`
      - ``SDF``. Collisions between (``SPHERE``, ``BOX``, ``MESH``, ``HFIELD``) and ``CYLINDER``. Collisions between
        (``BOX``, ``MESH``, ``HFIELD``) and ``ELLIPSOID``.
-   * - :ref:`Constraint <mjtConstraint>`
-     - :ref:`Frictionloss <coFriction>`, ``FRICTION_DOF``
    * - :ref:`Integrator <mjtIntegrator>`
      - ``IMPLICIT``
    * - Dynamics

mjx/mujoco/mjx/_src/constraint.py

@@ -45,6 +45,7 @@ class _Efc(PyTreeNode):
   solref: jax.Array
   solimp: jax.Array
   margin: jax.Array
+  frictionloss: jax.Array
 
 
 def _kbi(
@@ -119,8 +120,9 @@ def _efc_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
     j = (jacp1 - jacp2).T
     pos_imp = math.norm(pos)
     invweight = m.body_invweight0[obj1id, 0] + m.body_invweight0[obj2id, 0]
+    zero = jp.zeros_like(pos)
 
-    return _row(j, pos, pos_imp, invweight, solref, solimp, jp.zeros_like(pos))
+    return _row(j, pos, pos_imp, invweight, solref, solimp, zero, zero)
 
   args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp)
   args = jax.tree_util.tree_map(lambda x: x[eq_id], args)
@@ -165,8 +167,9 @@ def _efc_equality_weld(m: Model, d: Data) -> Optional[_Efc]:
     pos_imp = math.norm(pos)
     invweight = m.body_invweight0[obj1id] + m.body_invweight0[obj2id]
     invweight = jp.repeat(invweight, 3, axis=0)
+    zero = jp.zeros_like(pos)
 
-    return _row(j, pos, pos_imp, invweight, solref, solimp, jp.zeros_like(pos))
+    return _row(j, pos, pos_imp, invweight, solref, solimp, zero, zero)
 
   args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp)
   args = jax.tree_util.tree_map(lambda x: x[eq_id], args)
@@ -194,8 +197,9 @@ def _efc_equality_joint(m: Model, d: Data) -> Optional[_Efc]:
     j = jp.zeros((m.nv)).at[dofadr2].set(-deriv).at[dofadr1].set(1.0)
     invweight = m.dof_invweight0[dofadr1]
     invweight += m.dof_invweight0[dofadr2] * (obj2id > -1)
+    zero = jp.zeros_like(pos)
 
-    return _row(j, pos, pos, invweight, solref, solimp, jp.zeros_like(pos))
+    return _row(j, pos, pos, invweight, solref, solimp, zero, zero)
 
   args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp)
   args = jax.tree_util.tree_map(lambda x: x[eq_id], args)
@@ -232,8 +236,9 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
     pos = pos1 - jp.dot(data[:5], dif_power)
     deriv = jp.dot(data[1:5], dif_power[:4] * jp.arange(1, 5)) * (obj2id > -1)
     j = jac1 + jac2 * -deriv
+    zero = jp.zeros_like(pos)
 
-    return _row(j, pos, pos, invweight, solref, solimp, jp.zeros_like(pos))
+    return _row(j, pos, pos, invweight, solref, solimp, zero, zero)
 
   inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id]
   jac1, jac2 = d.ten_J[obj1id], d.ten_J[obj2id]
@@ -245,9 +250,32 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
 
 
 def _efc_friction(m: Model, d: Data) -> Optional[_Efc]:
-  # TODO(robotics-team): implement _instantiate_friction
+  """Calculates constraint rows for dof frictionloss."""
-  del m, d
+  dof_id = np.nonzero(m.dof_hasfrictionloss)[0]
-  return None
+  tendon_id = np.nonzero(m.tendon_hasfrictionloss)[0]
+
+  size = dof_id.size + tendon_id.size
+  if (m.opt.disableflags & DisableBit.FRICTIONLOSS) or (size == 0):
+    return None
+
+  args_dof = (jp.eye(m.nv), m.dof_frictionloss, m.dof_invweight0, m.dof_solref)
+  args_dof += (m.dof_solimp,)
+  args_dof = jax.tree_util.tree_map(lambda x: x[dof_id], args_dof)
+
+  args_ten = (d.ten_J, m.tendon_frictionloss, m.tendon_invweight0)
+  args_ten += (m.tendon_solref_fri, m.tendon_solimp_fri)
+  args_ten = jax.tree_util.tree_map(lambda x: x[tendon_id], args_ten)
+
+  args = jax.tree_util.tree_map(
+      lambda *x: jp.concatenate(x), args_dof, args_ten
+  )
+
+  @jax.vmap
+  def rows(j, frictionloss, invweight, solref, solimp):
+    z = jp.zeros_like(frictionloss)
+    return _row(j, z, z, invweight, solref, solimp, z, frictionloss)
+
+  return rows(*args)
 
 
 def _efc_limit_ball(m: Model, d: Data) -> Optional[_Efc]:
@@ -267,9 +295,10 @@ def _efc_limit_ball(m: Model, d: Data) -> Optional[_Efc]:
     active = pos < 0
     j = jp.zeros(m.nv).at[jp.arange(3) + dofadr].set(-axis)
     invweight = m.dof_invweight0[dofadr]
+    z = jp.zeros_like(pos)
 
     return _row(
-        j * active, pos * active, pos, invweight, solref, solimp, jnt_margin
+        j * active, pos * active, pos, invweight, solref, solimp, jnt_margin, z
     )
 
   args = (m.jnt_qposadr, m.jnt_dofadr, m.jnt_range, m.jnt_margin, m.jnt_solref)
@@ -296,9 +325,10 @@ def _efc_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
     active = pos < 0
     j = jp.zeros(m.nv).at[dofadr].set((dist_min < dist_max) * 2 - 1)
     invweight = m.dof_invweight0[dofadr]
+    z = jp.zeros_like(pos)
 
     return _row(
-        j * active, pos * active, pos, invweight, solref, solimp, jnt_margin
+        j * active, pos * active, pos, invweight, solref, solimp, jnt_margin, z
     )
 
   args = (m.jnt_qposadr, m.jnt_dofadr, m.jnt_range, m.jnt_margin, m.jnt_solref)
@@ -332,8 +362,11 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
   pos = jp.minimum(dist_min, dist_max) - margin
   active = pos < 0
   j = jax.vmap(jp.multiply)(j, ((dist_min < dist_max) * 2 - 1) * active)
+  zero = jp.zeros_like(pos)
 
-  return jax.vmap(_row)(j, pos * active, pos, invweight, solref, solimp, margin)
+  return jax.vmap(_row)(
+      j, pos * active, pos, invweight, solref, solimp, margin, zero
+  )
 
 
 def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
@@ -362,6 +395,7 @@ def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
         c.solref,
         c.solimp,
         c.includemargin,
+        jp.zeros_like(pos),
     )
 
   contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
@@ -406,6 +440,7 @@ def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]:
         c.solref,
         c.solimp,
         c.includemargin,
+        jp.zeros_like(pos),
     )
 
   contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
@@ -450,6 +485,7 @@ def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
         solref,
         c.solimp,
         c.includemargin,
+        jp.zeros_like(pos),
     )
 
   contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
@@ -460,7 +496,8 @@ def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
 def counts(efc_type: np.ndarray) -> Tuple[int, int, int, int]:
   """Returns equality, friction, limit, and contact constraint counts."""
   ne = (efc_type == ConstraintType.EQUALITY).sum()
-  nf = 0  # no support for friction loss yet
+  nf = (efc_type == ConstraintType.FRICTION_DOF).sum()
+  nf += (efc_type == ConstraintType.FRICTION_TENDON).sum()
   nl = (efc_type == ConstraintType.LIMIT_JOINT).sum()
   nl += (efc_type == ConstraintType.LIMIT_TENDON).sum()
   nc_f = (efc_type == ConstraintType.CONTACT_FRICTIONLESS).sum()
@@ -488,6 +525,20 @@ def make_efc_type(
     num_rows += (m.eq_type == EqType.TENDON).sum()
     efc_types += [ConstraintType.EQUALITY] * num_rows
 
+  if not m.opt.disableflags & DisableBit.FRICTIONLOSS:
+    nf_dof = (
+        m.dof_hasfrictionloss.sum()
+        if isinstance(m, Model)
+        else (m.dof_frictionloss > 0).sum()
+    )
+    efc_types += [ConstraintType.FRICTION_DOF] * nf_dof
+    nf_tendon = (
+        m.tendon_hasfrictionloss.sum()
+        if isinstance(m, Model)
+        else (m.tendon_frictionloss > 0).sum()
+    )
+    efc_types += [ConstraintType.FRICTION_TENDON] * nf_tendon
+
   if not m.opt.disableflags & DisableBit.LIMIT:
     efc_types += [ConstraintType.LIMIT_JOINT] * m.jnt_limited.sum()
     efc_types += [ConstraintType.LIMIT_TENDON] * m.tendon_limited.sum()
@@ -570,12 +621,12 @@ def make_constraint(m: Model, d: Data) -> Data:
     k, b, imp = _kbi(m, efc.solref, efc.solimp, efc.pos_imp)
     r = jp.maximum(efc.invweight * (1 - imp) / imp, mujoco.mjMINVAL)
     aref = -b * (efc.J @ d.qvel) - k * imp * efc.pos_aref
-    return aref, r, efc.pos_aref + efc.margin, efc.margin
+    return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss
 
-  aref, r, pos, margin = fn(efc)
+  aref, r, pos, margin, frictionloss = fn(efc)
   d = d.replace(
       efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos, efc_margin=margin
   )
-  d = d.replace(efc_frictionloss=jp.zeros_like(r))
+  d = d.replace(efc_frictionloss=frictionloss)
 
   return d

mjx/mujoco/mjx/_src/constraint_test.py

@@ -60,12 +60,17 @@ class ConstraintTest(parameterized.TestCase):
 
       order = test_util.efc_order(m, d, dx)
       d_efc_j = d.efc_J.reshape((-1, m.nv))
-      _assert_eq(d_efc_j, dx.efc_J[order][:d.nefc], 'efc_J')
+      _assert_eq(d_efc_j, dx.efc_J[order][: d.nefc], 'efc_J')
-      _assert_eq(0, dx.efc_J[order][d.nefc:], 'efc_J')
+      _assert_eq(0, dx.efc_J[order][d.nefc :], 'efc_J')
-      _assert_eq(d.efc_aref, dx.efc_aref[order][:d.nefc], 'efc_aref')
+      _assert_eq(d.efc_aref, dx.efc_aref[order][: d.nefc], 'efc_aref')
-      _assert_eq(0, dx.efc_aref[order][d.nefc:], 'efc_aref')
+      _assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref')
-      _assert_eq(d.efc_D, dx.efc_D[order][:d.nefc], 'efc_D')
+      _assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D')
-      _assert_eq(d.efc_pos, dx.efc_pos[order][:d.nefc], 'efc_pos')
+      _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
+      _assert_eq(
+          d.efc_frictionloss,
+          dx.efc_frictionloss[order][: d.nefc],
+          'efc_frictionloss',
+      )
 
   def test_disable_refsafe(self):
     m = test_util.load_test_file('constraints.xml')
@@ -96,22 +101,33 @@ class ConstraintTest(parameterized.TestCase):
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY
     ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
     self.assertEqual(ne, 0)
-    self.assertEqual(nf, 0)
+    self.assertEqual(nf, 2)
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 148)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 153)  # only joint/tendon limit, contact
+    self.assertEqual(dx.efc_J.shape[0], 155)  # only joint/tendon limit, contact
 
   def test_disable_contact(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
     ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
     self.assertEqual(ne, 11)
-    self.assertEqual(nf, 0)
+    self.assertEqual(nf, 2)
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 0)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 16)  # only equality, joint/tendon limit
+    self.assertEqual(dx.efc_J.shape[0], 18)  # only equality, joint/tendon limit
+
+  def test_disable_frictionloss(self):
+    m = test_util.load_test_file('constraints.xml')
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.FRICTIONLOSS
+    ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
+    self.assertEqual(ne, 11)
+    self.assertEqual(nf, 0)
+    self.assertEqual(nl, 5)
+    self.assertEqual(nc, 148)
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
+    self.assertEqual(dx.efc_J.shape[0], 164)
 
   def test_margin(self):
     """Test margin."""

mjx/mujoco/mjx/_src/io.py

@@ -126,15 +126,14 @@ def put_model(
           f'{[mj_type(m) for m in missing]} not supported'
       )
 
-  if not np.allclose(m.dof_frictionloss, 0) and not _full_compat:
-    raise NotImplementedError('dof_frictionloss is not implemented.')
-
   mj_field_names = {
       f.name
       for f in types.Model.fields()
       if f.metadata.get('restricted_to') != 'mjx'
   }
   fields = {f: getattr(m, f) for f in mj_field_names}
+  fields['dof_hasfrictionloss'] = fields['dof_frictionloss'] > 0
+  fields['tendon_hasfrictionloss'] = fields['tendon_frictionloss'] > 0
   fields['geom_rbound_hfield'] = fields['geom_rbound']
   fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
   fields['opt'] = _make_option(m.opt)

mjx/mujoco/mjx/_src/solver.py

@@ -145,9 +145,23 @@ class _LSPoint(PyTreeNode):
     """Creates a linesearch point with first and second derivatives."""
     # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
 
-    # TODO(robotics-team): change this to support friction constraints
     cost, deriv_0, deriv_1 = 0.0, 0.0, 0.0
     quad_total = quad_gauss
+    x = ctx.Jaref + alpha * jv
+    active = (x < 0).at[: d.ne + d.nf].set(True)
+
+    dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
+    if (dof_fl.any() or ten_fl.any()) and not (
+        m.opt.disableflags & DisableBit.FRICTIONLOSS
+    ):
+      f = d.efc_frictionloss
+      r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
+      rf, z = r * f, jp.zeros_like(f)
+      linear_neg = (x <= -rf)[:, None]
+      linear_pos = (x >= rf)[:, None]
+      qf = linear_neg * jp.array([f * (-0.5 * rf - ctx.Jaref), -f * jv, z]).T
+      qf += linear_pos * jp.array([f * (-0.5 * rf + ctx.Jaref), f * jv, z]).T
+      quad = jp.where(f[:, None] > 0, qf, quad)
 
     if m.opt.cone == ConeType.ELLIPTIC:
       mu, u0 = ctx.fri[:, 0], ctx.u[:, 0]
@@ -161,7 +175,6 @@ class _LSPoint(PyTreeNode):
       # quadratic cost for equality, friction, limits, frictionless contacts
       dim1 = d.contact.efc_address[d.contact.dim == 1]
       nefl = d.ne + d.nf + d.nl
-      active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
       active = active.at[nefl:].set(False).at[dim1].set(active[dim1])
       quad_efld = jax.vmap(jp.multiply)(quad, active)
       quad_total += jp.sum(quad_efld, axis=0)
@@ -181,7 +194,6 @@ class _LSPoint(PyTreeNode):
       deriv_0 = jp.sum(dm * nmt * (n1 - mu * t1))
       deriv_1 = jp.sum(dm * (jp.square(n1 - mu * t1) - nmt * mu * t2))
     elif m.opt.cone == ConeType.PYRAMIDAL:
-      active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
       quad = jax.vmap(jp.multiply)(quad, active)  # only active
       quad_total += jp.sum(quad, axis=0)
     else:
@@ -228,7 +240,7 @@ def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
 
 
 def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
-  """Updates constraint force and resulting cost given latst solver iteration.
+  """Updates constraint force and resulting cost given last solver iteration.
 
   Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c
 
@@ -240,10 +252,27 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   Returns:
     context with new constraint force and costs
   """
+  # ne constraints are always active, nf are conditionally active, others are
+  # non-negative constraints.
+  active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True)
+
+  floss_force, floss_cost = jp.zeros(d.nefc), 0.0
+  dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
+  if (dof_fl.any() or ten_fl.any()) and not (
+      m.opt.disableflags & DisableBit.FRICTIONLOSS
+  ):
+    f = d.efc_frictionloss
+    r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
+    linear_neg = (ctx.Jaref <= -r * f) * (f > 0)
+    linear_pos = (ctx.Jaref >= r * f) * (f > 0)
+    active = active & ~linear_neg & ~linear_pos
+    floss_force = linear_neg * f + linear_pos * -f
+    floss_cost = linear_neg * (-0.5 * r * f * f - f * ctx.Jaref)
+    floss_cost += linear_pos * (-0.5 * r * f * f + f * ctx.Jaref)
+    floss_cost = floss_cost.sum()
+
   if m.opt.cone == ConeType.PYRAMIDAL:
-    # ne/nf constraints are always active, rest are non-negative constraints
+    efc_force = d.efc_D * -ctx.Jaref * active + floss_force
-    active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
-    efc_force = d.efc_D * -ctx.Jaref * active
     cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
     dm, u, h = 0.0, 0.0, 0.0
   elif m.opt.cone == ConeType.ELLIPTIC:
@@ -256,13 +285,12 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
 
     # bottom zone: quadratic
     bottom_zone = ((t <= 0) & (n < 0)) | ((t > 0) & ((mu * n + t) <= 0))
-    active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
     adr_i, adr_j = [], []
     for i, (condim, addr) in enumerate(zip(dim, efc_address)):
       adr_i.extend(range(addr, addr + condim))
       adr_j.extend([i] * condim)
     active = active.at[jp.array(adr_i)].set(bottom_zone[jp.array(adr_j)])
-    efc_force = d.efc_D * -ctx.Jaref * active
+    efc_force = d.efc_D * -ctx.Jaref * active + floss_force
     cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
 
     # middle zone: cone
@@ -309,7 +337,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   ctx = ctx.replace(
       qfrc_constraint=qfrc_constraint,
       gauss=gauss,
-      cost=cost + gauss,
+      cost=cost + gauss + floss_cost,
       prev_cost=ctx.cost,
       efc_force=efc_force,
       active=active,

mjx/mujoco/mjx/_src/types.py

@@ -49,6 +49,7 @@ class DisableBit(enum.IntFlag):
   """
   CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
   EQUALITY = mujoco.mjtDisableBit.mjDSBL_EQUALITY
+  FRICTIONLOSS = mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS
   LIMIT = mujoco.mjtDisableBit.mjDSBL_LIMIT
   CONTACT = mujoco.mjtDisableBit.mjDSBL_CONTACT
   PASSIVE = mujoco.mjtDisableBit.mjDSBL_PASSIVE
@@ -60,7 +61,7 @@ class DisableBit(enum.IntFlag):
   SENSOR = mujoco.mjtDisableBit.mjDSBL_SENSOR
   EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP
   FILTERPARENT = mujoco.mjtDisableBit.mjDSBL_FILTERPARENT
-  # unsupported: FRICTIONLOSS, MIDPHASE
+  # unsupported: MIDPHASE
 
 
 class JointType(enum.IntEnum):
@@ -265,7 +266,8 @@ class ConstraintType(enum.IntEnum):
     CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
   """
   EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
-  # unsupported: FRICTION_DOF, FRICTION_TENDON
+  FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF
+  FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON
   LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT
   LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON
   CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS
@@ -557,6 +559,7 @@ class Model(PyTreeNode):
     dof_solref: constraint solver reference:frictionloss      (nv, mjNREF)
     dof_solimp: constraint solver impedance:frictionloss      (nv, mjNIMP)
     dof_frictionloss: dof friction loss                       (nv,)
+    dof_hasfrictionloss: dof has >0 frictionloss (MJX)        (nv,)
     dof_armature: dof armature inertia/mass                   (nv,)
     dof_damping: damping coefficient                          (nv,)
     dof_invweight0: diag. inverse inertia in qpos0            (nv,)
@@ -697,6 +700,7 @@ class Model(PyTreeNode):
     tendon_lengthspring: spring resting length range          (ntendon, 2)
     tendon_length0: tendon length in qpos0                    (ntendon,)
     tendon_invweight0: inv. weight in qpos0                   (ntendon,)
+    tendon_hasfrictionloss: tendon has >0 frictionloss (MJX)  (ntendon,)
     wrap_type: wrap object type (mjtWrap)                     (nwrap,)
     wrap_objid: object id: geom, site, joint                  (nwrap,)
     wrap_prm: divisor, joint coef, or site id                 (nwrap,)
@@ -860,6 +864,7 @@ class Model(PyTreeNode):
   dof_solref: jax.Array
   dof_solimp: jax.Array
   dof_frictionloss: jax.Array
+  dof_hasfrictionloss: np.ndarray = _restricted_to('mjx')
   dof_armature: jax.Array
   dof_damping: jax.Array
   dof_invweight0: jax.Array
@@ -1005,6 +1010,7 @@ class Model(PyTreeNode):
   tendon_lengthspring: jax.Array
   tendon_length0: jax.Array
   tendon_invweight0: jax.Array
+  tendon_hasfrictionloss: np.ndarray = _restricted_to('mjx')
   wrap_type: np.ndarray = _restricted_to('mujoco')
   wrap_objid: np.ndarray = _restricted_to('mujoco')
   wrap_prm: np.ndarray = _restricted_to('mujoco')

mjx/mujoco/mjx/test_data/constraints.xml

@@ -29,7 +29,7 @@
     </body>
 
     <body name="beam3" pos="1 0 0">
-      <joint name="joint3" axis="1 0 0" type="hinge" range="-20 20"/>
+      <joint name="joint3" axis="1 0 0" type="hinge" range="-20 20" frictionloss="0.1"/>
       <geom class="box"/>
     </body>
 
@@ -65,7 +65,7 @@
   </worldbody>
 
   <tendon>
-    <fixed name="tendon_1" limited="true" range="-0.3 0.1" stiffness=".1" damping=".2">
+    <fixed name="tendon_1" limited="true" range="-0.3 0.1" stiffness=".1" damping=".2" frictionloss="0.1">
       <joint joint="joint3" coef=".1"/>
       <joint joint="joint4" coef="-.2"/>
     </fixed>