From 49711fa14f82874d7a4f35e72d3bd01f8476dbda Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 3 Sep 2024 13:16:27 -0700 Subject: [PATCH] Add support for joint and tendon frictionloss to MJX. PiperOrigin-RevId: 670679777 Change-Id: I2eb56d013c11d4e1851f1c2edd6d98faaf854db9 --- doc/changelog.rst | 1 + doc/mjx.rst | 4 +- mjx/mujoco/mjx/_src/constraint.py | 79 +++++++++++++++++++----- mjx/mujoco/mjx/_src/constraint_test.py | 36 ++++++++--- mjx/mujoco/mjx/_src/io.py | 5 +- mjx/mujoco/mjx/_src/solver.py | 48 +++++++++++--- mjx/mujoco/mjx/_src/types.py | 10 ++- mjx/mujoco/mjx/test_data/constraints.xml | 4 +- 8 files changed, 143 insertions(+), 44 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 113eae38..338ed5fd 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -45,6 +45,7 @@ MJX - Added support for :ref:`implicitfast integration` for all cases except :doc:`fluid drag `. - Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices. +- Added support for joint and tendon :ref:`frictionloss `. Bug fixes ^^^^^^^^^ diff --git a/doc/mjx.rst b/doc/mjx.rst index 7ef9ca59..d5b3722a 100644 --- a/doc/mjx.rst +++ b/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 ` - - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL``, ``CONTACT_ELLIPTIC`` + - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``, ``FRICTION_TENDON`` * - :ref:`Equality ` - ``CONNECT``, ``WELD``, ``JOINT``, ``TENDON`` * - :ref:`Integrator ` @@ -228,8 +228,6 @@ The following features are **in development** and coming soon: * - :ref:`Geom ` - ``SDF``. Collisions between (``SPHERE``, ``BOX``, ``MESH``, ``HFIELD``) and ``CYLINDER``. Collisions between (``BOX``, ``MESH``, ``HFIELD``) and ``ELLIPSOID``. - * - :ref:`Constraint ` - - :ref:`Frictionloss `, ``FRICTION_DOF`` * - :ref:`Integrator ` - ``IMPLICIT`` * - Dynamics diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 40111298..4c317539 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/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 - del m, d - return None + """Calculates constraint rows for dof frictionloss.""" + dof_id = np.nonzero(m.dof_hasfrictionloss)[0] + 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 diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index b04ec6a1..8951d76f 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/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(0, dx.efc_J[order][d.nefc:], 'efc_J') - _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(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_j, 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(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_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.""" diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index eb059f9f..551cdc64 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index 31004f44..a288c833 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/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 - active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True) - 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) 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, diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 7dbe68dc..cfa777db 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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') diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml index a1476a45..49895470 100644 --- a/mjx/mujoco/mjx/test_data/constraints.xml +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -29,7 +29,7 @@ - + @@ -65,7 +65,7 @@ - +