diff --git a/doc/changelog.rst b/doc/changelog.rst index bb2a644a..b19bcc28 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -18,10 +18,14 @@ MJX ^^^ 4. Added more fields to ``mjx.Model`` and ``mjx.Data`` for further compatibility with the corresponding MuJoCo structs. +5. Added support for :ref:`fixed tendons `. +6. Added support for tendon length limits (``mjCNSTR_LIMIT_TENDON`` in :ref:`mjtConstraint`). +7. Added support for tendon equality constraints (``mjEQ_TENDON`` in :ref:`mjtEq`). +8. Added support for tendon actuator transmission (``mjTRN_TENDON`` in :ref:`mjtTrn`). Python bindings ^^^^^^^^^^^^^^^ -5. Added support for asset dictionary argument in ``mujoco.spec.from_file``, ``mujoco.spec.from_string`` and +9. Added support for asset dictionary argument in ``mujoco.spec.from_file``, ``mujoco.spec.from_string`` and ``mujoco.spec.compile``. Version 3.2.0 (Jul 15, 2024) diff --git a/doc/mjx.rst b/doc/mjx.rst index ec75e57c..1d8566be 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -188,7 +188,7 @@ The following features are **fully supported** in MJX: * - :ref:`Joint ` - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE`` * - :ref:`Transmission ` - - ``TRN_JOINT``, ``TRN_SITE`` + - ``TRN_JOINT``, ``TRN_SITE``, ``TRN_TENDON`` * - :ref:`Actuator Dynamics ` - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT`` * - :ref:`Actuator Gain ` @@ -200,7 +200,7 @@ The following features are **fully supported** in MJX: * - :ref:`Constraint ` - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL``, ``CONTACT_ELLIPTIC`` * - :ref:`Equality ` - - ``CONNECT``, ``WELD``, ``JOINT`` + - ``CONNECT``, ``WELD``, ``JOINT``, ``TENDON`` * - :ref:`Integrator ` - ``EULER``, ``RK4`` * - :ref:`Cone ` @@ -211,6 +211,8 @@ The following features are **fully supported** in MJX: - ``CG``, ``NEWTON`` * - Fluid Model - :ref:`flInertia` + * - :ref:`Tendons ` + - :ref:`Fixed ` The following features are **in development** and coming soon: @@ -230,8 +232,6 @@ The following features are **in development** and coming soon: - ``IMPLICIT``, ``IMPLICITFAST`` * - Dynamics - :ref:`Inverse ` - * - :ref:`Transmission ` - - ``TRN_TENDON`` * - :ref:`Actuator Dynamics ` - ``MUSCLE`` * - :ref:`Actuator Gain ` @@ -243,9 +243,7 @@ The following features are **in development** and coming soon: * - Fluid Model - :ref:`flEllipsoid` * - :ref:`Tendons ` - - :ref:`Spatial `, :ref:`Fixed ` - * - :ref:`Equality ` - - ``TENDON`` + - :ref:`Spatial ` * - :ref:`Sensors ` - All except ``PLUGIN``, ``USER`` * - Lights diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 0d59c95d..e0ad5cbe 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -39,6 +39,7 @@ from mujoco.mjx._src.smooth import crb from mujoco.mjx._src.smooth import factor_m from mujoco.mjx._src.smooth import kinematics from mujoco.mjx._src.smooth import rne +from mujoco.mjx._src.smooth import tendon from mujoco.mjx._src.smooth import transmission from mujoco.mjx._src.solver import solve from mujoco.mjx._src.support import full_m diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 5a1286d1..eb16863d 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -205,6 +205,44 @@ def _efc_equality_joint(m: Model, d: Data) -> Optional[_Efc]: return rows(*args) +def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]: + """Calculates constraint rows for tendon equality constraints.""" + + eq_id = np.nonzero(m.eq_type == EqType.TENDON)[0] + + if (m.opt.disableflags & DisableBit.EQUALITY) or eq_id.size == 0: + return None + + obj1id, obj2id, data, solref, solimp = jax.tree_util.tree_map( + lambda x: x[eq_id], + ( + m.eq_obj1id, + m.eq_obj2id, + m.eq_data, + m.eq_solref, + m.eq_solimp, + ), + ) + + @jax.vmap + def rows(obj2id, data, solref, solimp, invweight, jac1, jac2, pos1, pos2): + dif = pos2 * (obj2id > -1) + dif_power = jp.power(dif, jp.arange(0, 5)) + 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 + + return _row(j, pos, pos, invweight, solref, solimp) + + inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id] + jac1, jac2 = d.ten_J[obj1id], d.ten_J[obj2id] + pos1 = d.ten_length[obj1id] - m.tendon_length0[obj1id] + pos2 = d.ten_length[obj2id] - m.tendon_length0[obj2id] + invweight = inv1 + inv2 * (obj2id > -1) + + return rows(obj2id, data, solref, solimp, invweight, jac1, jac2, pos1, pos2) + + def _efc_friction(m: Model, d: Data) -> Optional[_Efc]: # TODO(robotics-team): implement _instantiate_friction del m, d @@ -222,6 +260,8 @@ def _efc_limit_ball(m: Model, d: Data) -> Optional[_Efc]: @jax.vmap def rows(qposadr, dofadr, jnt_range, jnt_margin, solref, solimp): axis, angle = math.quat_to_axis_angle(d.qpos[jp.arange(4) + qposadr]) + # ball rotation angle is always positive + axis, angle = math.normalize_with_norm(axis * angle) pos = jp.amax(jnt_range) - angle - jnt_margin active = pos < 0 j = jp.zeros(m.nv).at[jp.arange(3) + dofadr].set(-axis) @@ -263,6 +303,34 @@ def _efc_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]: return rows(*args) +def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]: + """Calculates constraint rows for tendon limits.""" + tendon_id = np.nonzero(m.tendon_limited)[0] + + if (m.opt.disableflags & DisableBit.LIMIT) or tendon_id.size == 0: + return None + + length, j, range_, margin, invweight, solref, solimp = jax.tree_util.tree_map( + lambda x: x[tendon_id], + ( + d.ten_length, + d.ten_J, + m.tendon_range, + m.tendon_margin, + m.tendon_invweight0, + m.tendon_solref_lim, + m.tendon_solimp_lim, + ), + ) + + dist_min, dist_max = length - range_[:, 0], range_[:, 1] - length + pos = jp.minimum(dist_min, dist_max) - margin + active = pos < 0 + j = jax.vmap(jp.multiply)(j, ((dist_min < dist_max) * 2 - 1) * active) + + return jax.vmap(_row)(j, pos * active, pos, invweight, solref, solimp) + + def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for frictionless contacts.""" @@ -365,6 +433,7 @@ def counts(efc_type: np.ndarray) -> Tuple[int, int, int, int]: ne = (efc_type == ConstraintType.EQUALITY).sum() nf = 0 # no support for friction loss yet nl = (efc_type == ConstraintType.LIMIT_JOINT).sum() + nl += (efc_type == ConstraintType.LIMIT_TENDON).sum() nc_f = (efc_type == ConstraintType.CONTACT_FRICTIONLESS).sum() nc_p = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum() nc_e = (efc_type == ConstraintType.CONTACT_ELLIPTIC).sum() @@ -387,10 +456,12 @@ def make_efc_type( num_rows = (m.eq_type == EqType.CONNECT).sum() * 3 num_rows += (m.eq_type == EqType.WELD).sum() * 6 num_rows += (m.eq_type == EqType.JOINT).sum() + num_rows += (m.eq_type == EqType.TENDON).sum() efc_types += [ConstraintType.EQUALITY] * num_rows 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() if not m.opt.disableflags & DisableBit.CONTACT: for condim in (1, 3, 4, 6): @@ -441,9 +512,11 @@ def make_constraint(m: Model, d: Data) -> Data: _efc_equality_connect(m, d), _efc_equality_weld(m, d), _efc_equality_joint(m, d), + _efc_equality_tendon(m, d), _efc_friction(m, d), _efc_limit_ball(m, d), _efc_limit_slide_hinge(m, d), + _efc_limit_tendon(m, d), _efc_contact_frictionless(m, d), ) if m.opt.cone == ConeType.ELLIPTIC: diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index f7cc6530..db13e699 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -96,21 +96,21 @@ class ConstraintTest(parameterized.TestCase): ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m)) self.assertEqual(ne, 0) self.assertEqual(nf, 0) - self.assertEqual(nl, 3) + 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], 151) # only joint range, contact + self.assertEqual(dx.efc_J.shape[0], 153) # 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, 10) + self.assertEqual(ne, 11) self.assertEqual(nf, 0) - self.assertEqual(nl, 3) + 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], 13) # only joint range, limit + self.assertEqual(dx.efc_J.shape[0], 16) # only equality, joint/tendon limit if __name__ == '__main__': diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index e2838d00..b161593c 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -66,6 +66,7 @@ def fwd_position(m: Model, d: Data) -> Data: d = smooth.kinematics(m, d) d = smooth.com_pos(m, d) d = smooth.camlight(m, d) + d = smooth.tendon(m, d) d = smooth.crb(m, d) d = smooth.factor_m(m, d) d = collision_driver.collision(m, d) @@ -77,7 +78,10 @@ def fwd_position(m: Model, d: Data) -> Data: @named_scope def fwd_velocity(m: Model, d: Data) -> Data: """Velocity-dependent computations.""" - d = d.replace(actuator_velocity=d.actuator_moment @ d.qvel) + d = d.replace( + actuator_velocity=d.actuator_moment @ d.qvel, + ten_velocity=d.ten_J @ d.qvel, + ) d = smooth.com_vel(m, d) d = passive.passive(m, d) d = smooth.rne(m, d) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index b9657c4f..ff6cd29e 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -15,7 +15,7 @@ """Functions to initialize, load, or save data.""" import copy -from typing import Any, Dict, List, Tuple, Union +from typing import List, Tuple, Union import jax from jax import numpy as jp @@ -74,9 +74,6 @@ def put_model( ) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" - if _check_unsupported and m.ntendon: - raise NotImplementedError('tendons are not supported') - mesh_geomid = set() for g1, g2, ip in collision_driver.geom_pairs(m): t1, t2 = m.geom_type[[g1, g2]] @@ -104,6 +101,7 @@ def put_model( (m.actuator_gaintype, types.GainType, mujoco.mjtGain), (m.actuator_trntype, types.TrnType, mujoco.mjtTrn), (m.eq_type, types.EqType, mujoco.mjtEq), + (m.wrap_type, types.WrapType, mujoco.mjtWrap), ): missing = set(enum_field) - set(enum_type) if _check_unsupported and missing: diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index bb540a5d..020e8e23 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -148,20 +148,28 @@ class ModelIOTest(parameterized.TestCase): ) ) - def test_tendon_not_implemented(self): + def test_spatial_tendon_not_implemented(self): with self.assertRaises(NotImplementedError): mjx.put_model(mujoco.MjModel.from_xml_string(""" - - - + + + + + + + + + + - - - + + + + """)) diff --git a/mjx/mujoco/mjx/_src/passive.py b/mjx/mujoco/mjx/_src/passive.py index c2428566..0975b129 100644 --- a/mjx/mujoco/mjx/_src/passive.py +++ b/mjx/mujoco/mjx/_src/passive.py @@ -70,6 +70,13 @@ def _spring_damper(m: Model, d: Data) -> jax.Array: # dof-level dampers qfrc -= m.dof_damping * d.qvel + # tendon-level spring-dampers + below, above = m.tendon_lengthspring.T - d.ten_length + frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0) + frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring) + frc_damper = -m.tendon_damping * d.ten_velocity + qfrc += d.ten_J.T @ (frc_spring + frc_damper) + return qfrc diff --git a/mjx/mujoco/mjx/_src/passive_test.py b/mjx/mujoco/mjx/_src/passive_test.py index 910ce4d3..30bc8c66 100644 --- a/mjx/mujoco/mjx/_src/passive_test.py +++ b/mjx/mujoco/mjx/_src/passive_test.py @@ -42,7 +42,7 @@ class PassiveTest(absltest.TestCase): m = test_util.load_test_file('pendula.xml') d = mujoco.MjData(m) # give the system a little kick to ensure we have non-identity rotations - d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4, 0.5, -0.6, 0.1]) + d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4, 0.5, -0.6, 0.1, -0.2]) mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero mujoco.mj_forward(m, d) mx = mjx.put_model(m) diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 7f105199..ed6f3dce 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -311,7 +311,9 @@ def factor_m(m: Model, d: Data) -> Data: if j == -1: break out_beg, out_end = tuple(m.dof_Madr[j : j + 2]) - updates.setdefault(depth[j], []).append((out_beg, out_end, madr_d, madr_ij)) + updates.setdefault(depth[j], []).append( + (out_beg, out_end, madr_d, madr_ij) + ) qld = d.qM @@ -463,6 +465,20 @@ def rne(m: Model, d: Data) -> Data: return d +def tendon(m: Model, d: Data) -> Data: + """Computes tendon lengths and moments.""" + if not m.ntendon: + return d + + ten_id = np.repeat(np.arange(m.ntendon), m.tendon_num) + length = m.wrap_prm * d.qpos[m.jnt_qposadr[m.wrap_objid]] + ten_length = jax.ops.segment_sum(length, ten_id, m.ntendon) + ten_j = jp.zeros((m.ntendon, m.nv)) + ten_j = ten_j.at[ten_id, m.jnt_dofadr[m.wrap_objid]].set(m.wrap_prm) + + return d.replace(ten_length=ten_length, ten_J=ten_j) + + def _site_dof_mask(m: Model) -> np.ndarray: """Creates a dof mask for site transmissions.""" mask = np.ones((m.nu, m.nv)) @@ -494,7 +510,6 @@ def _site_dof_mask(m: Model) -> np.ndarray: def transmission(m: Model, d: Data) -> Data: """Computes actuator/transmission lengths and moments.""" - # TODO: consider combining transmission calculation into fwd_actuation. if not m.nu: return d @@ -545,6 +560,9 @@ def transmission(m: Model, d: Data) -> Data: jac = jp.concatenate((jacp, jacr), axis=1) * site_dof_mask[:, None] wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:])) moment = jac @ wrench + elif trntype == TrnType.TENDON: + length = d.ten_length[trnid[0]] * gear[:1] + moment = d.ten_J[trnid[0]] * gear[0] else: raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}') diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 0b25f71f..f0efc07f 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -48,7 +48,7 @@ class SmoothTest(absltest.TestCase): """Tests MJX smooth functions match MuJoCo smooth functions.""" m = test_util.load_test_file('pendula.xml') - # # force MJX sparse for testing: + # tell MJX to use sparse mass matrices: m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE d = mujoco.MjData(m) # give the system a little kick to ensure we have non-identity rotations @@ -75,6 +75,10 @@ class SmoothTest(absltest.TestCase): _assert_attr_eq(d, dx, 'subtree_com') _assert_attr_eq(d, dx, 'cinert') _assert_attr_eq(d, dx, 'cdof') + # camlight + dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'cam_xpos') + _assert_eq(d.cam_xmat.reshape((-1, 3, 3)), dx.cam_xmat, 'cam_xmat') # crb dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d)) _assert_attr_eq(d, dx, 'crb') @@ -90,14 +94,22 @@ class SmoothTest(absltest.TestCase): # rne dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d)) _assert_attr_eq(d, dx, 'qfrc_bias') + + # set dense jacobian for tendon: + m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE + d = mujoco.MjData(m) + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mujoco.mj_forward(m, d) + # tendon + dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d)) + _assert_attr_eq(d, dx, 'ten_J') + _assert_attr_eq(d, dx, 'ten_length') # transmission - dx = jax.jit(mjx.transmission)(mx, mjx.put_data(m, d)) + dx = jax.jit(mjx.transmission)(mx, dx) _assert_attr_eq(d, dx, 'actuator_length') _assert_attr_eq(d, dx, 'actuator_moment') - # camlight - dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d)) - _assert_attr_eq(d, dx, 'cam_xpos') - _assert_eq(d.cam_xmat.reshape((-1, 3, 3)), dx.cam_xmat, 'cam_xmat') def test_disable_gravity(self): m = mujoco.MjModel.from_xml_string(""" diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index d9e3375c..0500eefe 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -25,7 +25,7 @@ import numpy as np class DisableBit(enum.IntFlag): """Disable default feature bitflags. - Attributes: + Members: CONSTRAINT: entire constraint solver EQUALITY: equality constraints FRICTIONLOSS: joint and tendon frictionloss constraints @@ -56,7 +56,7 @@ class DisableBit(enum.IntFlag): class JointType(enum.IntEnum): """Type of degree of freedom. - Attributes: + Members: FREE: global position and orientation (quat) (7,) BALL: orientation (quat) relative to parent (4,) SLIDE: sliding distance along body-fixed axis (1,) @@ -77,7 +77,7 @@ class JointType(enum.IntEnum): class IntegratorType(enum.IntEnum): """Integrator mode. - Attributes: + Members: EULER: semi-implicit Euler RK4: 4th-order Runge Kutta """ @@ -89,7 +89,7 @@ class IntegratorType(enum.IntEnum): class GeomType(enum.IntEnum): """Type of geometry. - Attributes: + Members: PLANE: plane HFIELD: height field SPHERE: sphere @@ -115,7 +115,7 @@ class GeomType(enum.IntEnum): class ConvexMesh(PyTreeNode): """Geom properties for convex meshes. - Attributes: + Members: vert: vertices of the convex mesh face: faces of the convex mesh face_normal: normal vectors for the faces @@ -133,7 +133,7 @@ class ConvexMesh(PyTreeNode): class ConeType(enum.IntEnum): """Type of friction cone. - Attributes: + Members: PYRAMIDAL: pyramidal ELLIPTIC: elliptic """ @@ -144,7 +144,7 @@ class ConeType(enum.IntEnum): class JacobianType(enum.IntEnum): """Type of constraint Jacobian. - Attributes: + Members: DENSE: dense SPARSE: sparse AUTO: sparse if nv>60 and device is TPU, dense otherwise @@ -157,7 +157,7 @@ class JacobianType(enum.IntEnum): class SolverType(enum.IntEnum): """Constraint solver algorithm. - Attributes: + Members: CG: Conjugate gradient (primal) """ # unsupported: PGS @@ -168,33 +168,46 @@ class SolverType(enum.IntEnum): class EqType(enum.IntEnum): """Type of equality constraint. - Attributes: + Members: CONNECT: connect two bodies at a point (ball joint) WELD: fix relative position and orientation of two bodies JOINT: couple the values of two scalar joints with cubic + TENDON: couple the lengths of two tendons with cubic """ CONNECT = mujoco.mjtEq.mjEQ_CONNECT WELD = mujoco.mjtEq.mjEQ_WELD JOINT = mujoco.mjtEq.mjEQ_JOINT - # unsupported: TENDON, DISTANCE + TENDON = mujoco.mjtEq.mjEQ_TENDON + # unsupported: DISTANCE + + +class WrapType(enum.IntEnum): + """Type of tendon wrap object. + + Members: + JOINT: constant moment arm + """ + JOINT = mujoco.mjtWrap.mjWRAP_JOINT + # unsupported: NONE, PULLEY, SITE, SPHERE, CYLINDER class TrnType(enum.IntEnum): """Type of actuator transmission. - Attributes: + Members: JOINT: force on joint SITE: force on site """ JOINT = mujoco.mjtTrn.mjTRN_JOINT SITE = mujoco.mjtTrn.mjTRN_SITE - # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, BODY + TENDON = mujoco.mjtTrn.mjTRN_TENDON + # unsupported: JOINTINPARENT, SLIDERCRANK, BODY class DynType(enum.IntEnum): """Type of actuator dynamics. - Attributes: + Members: NONE: no internal dynamics; ctrl specifies force INTEGRATOR: integrator: da/dt = u FILTER: linear filter: da/dt = (u-a) / tau @@ -210,7 +223,7 @@ class DynType(enum.IntEnum): class GainType(enum.IntEnum): """Type of actuator gain. - Attributes: + Members: FIXED: fixed gain AFFINE: const + kp*length + kv*velocity """ @@ -222,7 +235,7 @@ class GainType(enum.IntEnum): class BiasType(enum.IntEnum): """Type of actuator bias. - Attributes: + Members: NONE: no bias AFFINE: const + kp*length + kv*velocity """ @@ -234,16 +247,17 @@ class BiasType(enum.IntEnum): class ConstraintType(enum.IntEnum): """Type of constraint. - Attributes: + Members: EQUALITY: equality constraint LIMIT_JOINT: joint limit + LIMIT_TENDON: tendon limit CONTACT_FRICTIONLESS: frictionless contact CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone """ EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY # unsupported: FRICTION_DOF, FRICTION_TENDON LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT - # unsupported: LIMIT_TENDON + LIMIT_TENDON = mujoco.mjtConstraint.mjCNSTR_LIMIT_TENDON CONTACT_FRICTIONLESS = mujoco.mjtConstraint.mjCNSTR_CONTACT_FRICTIONLESS CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL CONTACT_ELLIPTIC = mujoco.mjtConstraint.mjCNSTR_CONTACT_ELLIPTIC @@ -252,7 +266,7 @@ class ConstraintType(enum.IntEnum): class CamLightType(enum.IntEnum): """Type of camera light. - Attributes: + Members: FIXED: pos and rot fixed in body TRACK: pos tracks body, rot fixed in global TRACKCOM: pos tracks subtree com, rot fixed in body @@ -900,18 +914,18 @@ class Model(PyTreeNode): tendon_adr: np.ndarray tendon_num: np.ndarray tendon_limited: np.ndarray - tendon_solref_lim: np.ndarray - tendon_solimp_lim: np.ndarray - tendon_solref_fri: np.ndarray - tendon_solimp_fri: np.ndarray - tendon_range: np.ndarray - tendon_margin: np.ndarray - tendon_stiffness: np.ndarray - tendon_damping: np.ndarray - tendon_frictionloss: np.ndarray - tendon_lengthspring: np.ndarray - tendon_length0: np.ndarray - tendon_invweight0: np.ndarray + tendon_solref_lim: jax.Array + tendon_solimp_lim: jax.Array + tendon_solref_fri: jax.Array + tendon_solimp_fri: jax.Array + tendon_range: jax.Array + tendon_margin: jax.Array + tendon_stiffness: jax.Array + tendon_damping: jax.Array + tendon_frictionloss: jax.Array + tendon_lengthspring: jax.Array + tendon_length0: jax.Array + tendon_invweight0: jax.Array wrap_type: np.ndarray wrap_objid: np.ndarray wrap_prm: np.ndarray @@ -1193,7 +1207,7 @@ class Data(PyTreeNode): qM: jax.Array # pylint:disable=invalid-name qLD: jax.Array # pylint:disable=invalid-name qLDiagInv: jax.Array # pylint:disable=invalid-name - qLDiagSqrtInv: jax.Array + qLDiagSqrtInv: jax.Array # pylint:disable=invalid-name bvh_aabb_dyn: jax.Array bvh_active: jax.Array # position, velocity dependent: @@ -1243,6 +1257,6 @@ class Data(PyTreeNode): efc_force: jax.Array # sparse representation of qM, qLD, qLDiagInv, for compatibility with MuJoCo # when in dense mode - _qM_sparse: jax.Array - _qLD_sparse: jax.Array - _qLDiagInv_sparse: jax.Array + _qM_sparse: jax.Array # pylint:disable=invalid-name + _qLD_sparse: jax.Array # pylint:disable=invalid-name + _qLDiagInv_sparse: jax.Array # pylint:disable=invalid-name diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml index f79ee1e2..8f7e5643 100644 --- a/mjx/mujoco/mjx/test_data/constraints.xml +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -64,10 +64,22 @@ + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/pendula.xml b/mjx/mujoco/mjx/test_data/pendula.xml index aff7741a..82d2f23f 100644 --- a/mjx/mujoco/mjx/test_data/pendula.xml +++ b/mjx/mujoco/mjx/test_data/pendula.xml @@ -126,6 +126,17 @@ + + + + + + + + + + + @@ -135,5 +146,6 @@ +