Cleanup MJX tests and migrate them to put_model/put_data.

PiperOrigin-RevId: 588304729
Change-Id: I58075383e6eed64ae00cea305a568eba6465b0b0

mjx/mujoco/mjx/__init__.py

@@ -16,10 +16,17 @@
 
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
+from mujoco.mjx._src.constraint import count_constraints
 from mujoco.mjx._src.constraint import make_constraint
 from mujoco.mjx._src.device import device_get_into
 from mujoco.mjx._src.device import device_put
+from mujoco.mjx._src.forward import euler
 from mujoco.mjx._src.forward import forward
+from mujoco.mjx._src.forward import fwd_acceleration
+from mujoco.mjx._src.forward import fwd_actuation
+from mujoco.mjx._src.forward import fwd_position
+from mujoco.mjx._src.forward import fwd_velocity
+from mujoco.mjx._src.forward import rungekutta4
 from mujoco.mjx._src.forward import step
 from mujoco.mjx._src.io import get_data
 from mujoco.mjx._src.io import make_data
@@ -34,4 +41,5 @@ from mujoco.mjx._src.smooth import kinematics
 from mujoco.mjx._src.smooth import mul_m
 from mujoco.mjx._src.smooth import rne
 from mujoco.mjx._src.smooth import transmission
+from mujoco.mjx._src.solver import solve
 from mujoco.mjx._src.types import *

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -52,9 +52,9 @@ def _collide(
     mjcf: str, assets: Optional[Dict[str, str]] = None
 ) -> Tuple[mujoco.MjModel, mujoco.MjData, Model, Data]:
   m = mujoco.MjModel.from_xml_string(mjcf, assets or {})
-  mx = mjx.device_put(m)
+  mx = mjx.put_model(m)
   d = mujoco.MjData(m)
-  dx = mjx.device_put(d)
+  dx = mjx.put_data(m, d)
 
   mujoco.mj_step(m, d)
   collision_jit_fn = jax.jit(mjx.collision)
@@ -418,9 +418,9 @@ class BodyPairFilterTest(absltest.TestCase):
   def test_filter_parent_child(self):
     """Tests that parent-child collisions get filtered."""
     m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_step(m, d)
     collision_jit_fn = jax.jit(mjx.collision)
@@ -435,9 +435,9 @@ class BodyPairFilterTest(absltest.TestCase):
     """Tests that filterparent flag disables parent-child filtering."""
     m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD)
     m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_FILTERPARENT
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_step(m, d)
     collision_jit_fn = jax.jit(mjx.collision)
@@ -454,22 +454,14 @@ class NconTest(parameterized.TestCase):
   """Tests ncon."""
 
   def test_ncon(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    ncon = collision_driver.ncon(m)
-    d.qpos[2] = 0.0
+    self.assertEqual(ncon, 16)
-
-    mx = mjx.device_put(m)
-    ncon = collision_driver.ncon(mx)
-    self.assertEqual(ncon, 4)
 
   def test_disable_contact(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    m.opt.disableflags |= DisableBit.CONTACT
-    d.qpos[2] = 0.0
+    ncon = collision_driver.ncon(m)
-
-    m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT
-    mx = mjx.device_put(m)
-    ncon = collision_driver.ncon(mx)
     self.assertEqual(ncon, 0)
 
 
@@ -500,12 +492,12 @@ class TopKContactTest(absltest.TestCase):
 
   def test_top_k_contacts(self):
     m = mujoco.MjModel.from_xml_string(self._CAPSULES)
-    mx_top_k = mjx.device_put(m)
+    mx_top_k = mjx.put_model(m)
     mx_all = mx_top_k.replace(
         nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([])
     )
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     collision_jit_fn = jax.jit(mjx.collision)
     kinematics_jit_fn = jax.jit(mjx.kinematics)

mjx/mujoco/mjx/_src/constraint_test.py

@@ -15,159 +15,92 @@
 """Tests for constraint functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
-import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import constraint
 from mujoco.mjx._src import test_util
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.types import DisableBit
-from mujoco.mjx._src.types import SolverType
-# pylint: enable=g-importing-member
 import numpy as np
 
 
-def _assert_eq(a, b, name, step, fname, atol=5e-3, rtol=5e-3):
+# tolerance for difference between MuJoCo and MJX constraint calculations,
-  err_msg = f'mismatch: {name} at step {step} in {fname}'
+# mostly due to float precision
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
+_TOLERANCE = 5e-5
 
 
-class ConstraintTest(parameterized.TestCase):
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10   # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  @parameterized.parameters(enumerate(test_util.TEST_FILES))
+
-  def test_constraints(self, seed, fname):
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class ConstraintTest(absltest.TestCase):
+
+  def test_constraints(self):
     """Test constraints."""
-    np.random.seed(seed)
+    m = test_util.load_test_file('constraints.xml')
-
-    # exclude convex.xml since convex contacts are not exactly equivalent
-    if fname == 'convex.xml':
-      return
-
-    m = test_util.load_test_file(fname)
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
-    dx = mjx.make_data(mx)
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
 
-    forward_jit_fn = jax.jit(mjx.forward)
+    dx = mjx.make_constraint(mx, dx)
-
+    nnz = dx.efc_J.any(axis=1)
-    # give the system a little kick to ensure we have non-identity rotations
+    _assert_eq(d.efc_J, dx.efc_J[nnz].reshape(-1), 'efc_J')
-    d.qvel = np.random.random(m.nv)
+    _assert_eq(d.efc_D, dx.efc_D[nnz], 'efc_D')
-    for i in range(100):
+    _assert_eq(d.efc_aref, dx.efc_aref[nnz], 'efc_aref')
-      dx = dx.replace(qpos=jax.device_put(d.qpos), qvel=jax.device_put(d.qvel))
+    _assert_eq(d.efc_frictionloss, dx.efc_frictionloss[nnz], 'efc_frictionloss')
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx)
-
-      nnz_filter = dx.efc_J.any(axis=1)
-
-      mj_efc_j = d.efc_J.reshape((-1, m.nv))
-      mjx_efc_j = dx.efc_J[nnz_filter]
-      _assert_eq(mj_efc_j, mjx_efc_j, 'efc_J', i, fname)
-
-      mjx_efc_d = dx.efc_D[nnz_filter]
-      _assert_eq(d.efc_D, mjx_efc_d, 'efc_D', i, fname)
-
-      mjx_efc_aref = dx.efc_aref[nnz_filter]
-      _assert_eq(d.efc_aref, mjx_efc_aref, 'efc_aref', i, fname)
-
-      mjx_efc_frictionloss = dx.efc_frictionloss[nnz_filter]
-      _assert_eq(
-          d.efc_frictionloss,
-          mjx_efc_frictionloss,
-          'efc_frictionloss',
-          i,
-          fname,
-      )
-
-  _JNT_RANGE = """
-    <mujoco>
-      <worldbody>
-        <body pos="0 0 1">
-          <joint type="slide" axis="1 0 0" range="-1.8 1.8" solreflimit=".08 1"
-           damping="5e-4"/>
-          <geom type="box" size="0.2 0.15 0.1" mass="1"/>
-          <body>
-            <joint axis="0 1 0" damping="2e-6"/>
-            <geom type="capsule" fromto="0 0 0 0 0 1" size="0.045" mass=".1"/>
-          </body>
-        </body>
-      </worldbody>
-    </mujoco>
-  """
-
-  def test_jnt_range(self):
-    """Tests that mixed joint ranges are respected."""
-    # TODO(robotics-simulation): also test ball
-    m = mujoco.MjModel.from_xml_string(self._JNT_RANGE)
-    m.opt.solver = SolverType.CG.value
-    d = mujoco.MjData(m)
-    d.qpos = np.array([2.0, 15.0])
-
-    mx = mjx.device_put(m)
-    dx = mjx.device_put(d)
-    efc = jax.jit(constraint._instantiate_limit_slide_hinge)(mx, dx)
-
-    # first joint is outside the joint range
-    np.testing.assert_array_almost_equal(efc.J[0, 0], -1.0)
-
-    # second joint has no range, so only one efc row
-    self.assertEqual(efc.J.shape[0], 1)
 
   def test_disable_refsafe(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
 
     timeconst = m.opt.timestep / 4.0  # timeconst < 2 * timestep
     solimp = jp.array([timeconst, 1.0])
     solref = jp.array([0.8, 0.99, 0.001, 0.2, 2])
     pos = jp.ones(3)
 
-    m.opt.disableflags = m.opt.disableflags | DisableBit.REFSAFE
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.REFSAFE
     mx = mjx.device_put(m)
     k, *_ = constraint._kbi(mx, solimp, solref, pos)
     self.assertEqual(k, 1 / (0.99**2 * timeconst**2))
 
-    m.opt.disableflags = m.opt.disableflags & ~DisableBit.REFSAFE
+  def test_disable_constraint(self):
-    mx = mjx.device_put(m)
+    m = test_util.load_test_file('constraints.xml')
-    k, *_ = constraint._kbi(mx, solimp, solref, pos)
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONSTRAINT
-    self.assertEqual(k, 1 / (0.99**2 * (2 * m.opt.timestep) ** 2))
+    ne, nf, nl, nc = mjx.count_constraints(m)
-
+    self.assertEqual(ne, 0)
-  def test_disableconstraint(self):
+    self.assertEqual(nf, 0)
-    m = test_util.load_test_file('ant.xml')
+    self.assertEqual(nl, 0)
-    d = mujoco.MjData(m)
+    self.assertEqual(nc, 0)
-
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
-    dx = constraint.make_constraint(mx, dx)
     self.assertEqual(dx.efc_J.shape[0], 0)
 
   def test_disable_equality(self):
-    m = test_util.load_test_file('equality.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY
-
+    ne, nf, nl, nc = mjx.count_constraints(m)
-    m.opt.disableflags = m.opt.disableflags | DisableBit.EQUALITY
+    self.assertEqual(ne, 0)
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
+    self.assertEqual(nf, 0)
-    dx = constraint.make_constraint(mx, dx)
+    self.assertEqual(nl, 2)
-    self.assertEqual(dx.efc_J.shape[0], 0)
+    self.assertEqual(nc, 64)
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
+    self.assertEqual(dx.efc_J.shape[0], 66)  # only joint range, contact
 
   def test_disable_contact(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
-    d.qpos[2] = 0.0
+    ne, nf, nl, nc = mjx.count_constraints(m)
-    mujoco.mj_forward(m, d)
+    self.assertEqual(ne, 10)
-
+    self.assertEqual(nf, 0)
-    m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONTACT
+    self.assertEqual(nl, 2)
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
+    self.assertEqual(nc, 0)
-    dx = dx.tree_replace(
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-        {'contact.frame': dx.contact.frame.reshape((-1, 3, 3))}
+    self.assertEqual(dx.efc_J.shape[0], 12)  # only joint range, limit
-    )
-    efc = constraint._instantiate_contact(mx, dx)
-    self.assertIsNotNone(efc)
-
-    m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
-    efc = constraint._instantiate_contact(mx, dx)
-    self.assertIsNone(efc)
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/device.py

@@ -184,6 +184,11 @@ def device_put(value):
   Returns:
     on-device MJX struct reflecting the input value
   """
+  warnings.warn(
+      'device_put is deprecated, use put_model and put_data instead',
+      category=DeprecationWarning,
+  )
+
   clz = _TYPE_MAP.get(type(value))
   if clz is None:
     raise NotImplementedError(f'{type(value)} is not supported for device_put.')
@@ -242,6 +247,10 @@ def device_get_into(result, value):
   Raises:
     RuntimeError: if result length doesn't match data batch size
   """
+  warnings.warn(
+      'device_get_into is deprecated, use get_data instead',
+      category=DeprecationWarning,
+  )
 
   value = jax.device_get(value)
 

mjx/mujoco/mjx/_src/device_test.py

@@ -130,31 +130,31 @@ class ValidateInputTest(absltest.TestCase):
       mjx.device_put(m)
 
   def test_trn(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     m.actuator_trntype[0] = mujoco.mjtTrn.mjTRN_SITE
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
   def test_dyn(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     m.actuator_dyntype[0] = mujoco.mjtDyn.mjDYN_MUSCLE
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
   def test_gain(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
   def test_bias(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
   def test_condim(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
     for i in [1, 4, 6]:
       m.geom_condim[0] = i
       with self.assertRaises(NotImplementedError):

mjx/mujoco/mjx/_src/forward.py

@@ -60,7 +60,7 @@ def named_scope(fn, name: str = ''):
 
 
 @named_scope
-def _position(m: Model, d: Data) -> Data:
+def fwd_position(m: Model, d: Data) -> Data:
   """Position-dependent computations."""
   # TODO(robotics-simulation): tendon
   d = smooth.kinematics(m, d)
@@ -74,7 +74,7 @@ def _position(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _velocity(m: Model, d: Data) -> Data:
+def fwd_velocity(m: Model, d: Data) -> Data:
   """Velocity-dependent computations."""
   d = d.replace(actuator_velocity=d.actuator_moment @ d.qvel)
   d = smooth.com_vel(m, d)
@@ -84,7 +84,7 @@ def _velocity(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _actuation(m: Model, d: Data) -> Data:
+def fwd_actuation(m: Model, d: Data) -> Data:
   """Actuation-dependent computations."""
   if not m.nu or m.opt.disableflags & DisableBit.ACTUATION:
     return d.replace(
@@ -190,7 +190,7 @@ def _actuation(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _acceleration(m: Model, d: Data) -> Data:
+def fwd_acceleration(m: Model, d: Data) -> Data:
   """Add up all non-constraint forces, compute qacc_smooth."""
   qfrc_applied = d.qfrc_applied + support.xfrc_accumulate(m, d)
   qfrc_smooth = d.qfrc_passive - d.qfrc_bias + d.qfrc_actuator + qfrc_applied
@@ -263,7 +263,7 @@ def _advance(
 
 
 @named_scope
-def _euler(m: Model, d: Data) -> Data:
+def euler(m: Model, d: Data) -> Data:
   """Euler integrator, semi-implicit in velocity."""
   # integrate damping implicitly
   qacc = d.qacc
@@ -277,7 +277,7 @@ def _euler(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _rungekutta4(m: Model, d: Data) -> Data:
+def rungekutta4(m: Model, d: Data) -> Data:
   """Runge-Kutta explicit order 4 integrator."""
   d_t0 = d
   # pylint: disable=invalid-name
@@ -323,10 +323,10 @@ def _rungekutta4(m: Model, d: Data) -> Data:
 @named_scope
 def forward(m: Model, d: Data) -> Data:
   """Forward dynamics."""
-  d = _position(m, d)
+  d = fwd_position(m, d)
-  d = _velocity(m, d)
+  d = fwd_velocity(m, d)
-  d = _actuation(m, d)
+  d = fwd_actuation(m, d)
-  d = _acceleration(m, d)
+  d = fwd_acceleration(m, d)
 
   if d.efc_J.size == 0:
     d = d.replace(qacc=d.qacc_smooth)
@@ -343,9 +343,9 @@ def step(m: Model, d: Data) -> Data:
   d = forward(m, d)
 
   if m.opt.integrator == IntegratorType.EULER:
-    d = _euler(m, d)
+    d = euler(m, d)
   elif m.opt.integrator == IntegratorType.RK4:
-    d = _rungekutta4(m, d)
+    d = rungekutta4(m, d)
   else:
     raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.')
 

mjx/mujoco/mjx/_src/forward_test.py

@@ -15,77 +15,75 @@
 """Tests for forward functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
 import jax
-from jax import numpy as jp
 import mujoco
 from mujoco import mjx
-from mujoco.mjx._src import forward
 from mujoco.mjx._src import test_util
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.types import DisableBit
-# pylint: enable=g-importing-member
 import numpy as np
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-3, rtol=1e-3):
+# tolerance for difference between MuJoCo and MJX forward calculations - mostly
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+# due to float precision
-  a, b = getattr(a, attr), getattr(b, attr)
+_TOLERANCE = 1e-5
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-class ForwardTest(parameterized.TestCase):
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  @parameterized.parameters(
-      filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES)
-  )
-  def test_forward(self, fname):
-    """Test mujoco mj forward function matches mujoco_mjx forward function."""
-    np.random.seed(test_util.TEST_FILES.index(fname))
 
-    m = test_util.load_test_file(fname)
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class ForwardTest(absltest.TestCase):
+
+  def test_forward(self):
+    m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    # apply some control and xfrc input
-    dx = mjx.make_data(mx)
+    d.ctrl = np.array([-18, 0.59, 0.47])
-    forward_jit_fn = jax.jit(mjx.forward)
+    d.xfrc_applied[0, 2] = 0.1  # torque
+    d.xfrc_applied[1, 4] = 0.3  # linear force
+    mujoco.mj_step(m, d, 100)  # get some dynamics going
+    mujoco.mj_forward(m, d)
 
-    # give the system a little kick to ensure we have non-identity rotations
+    mx = mjx.put_model(m)
-    d.qvel = np.random.random(m.nv) * 0.05
-    for i in range(100):
-      qpos, qvel = d.qpos.copy(), d.qvel.copy()
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
 
-      _assert_attr_eq(d, dx, 'qfrc_smooth', i, fname)
+    # fwd_actuation
-      _assert_attr_eq(d, dx, 'qacc_smooth', i, fname)
+    dx = jax.jit(mjx.fwd_actuation)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'act_dot')
+    _assert_attr_eq(d, dx, 'qfrc_actuator')
 
-  @parameterized.parameters(
+    # fwd_accleration (fwd_position and fwd_velocity already tested elsewhere)
-      filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES)
+    dx = jax.jit(mjx.fwd_acceleration)(mx, mjx.put_data(m, d))
-  )
+    _assert_attr_eq(d, dx, 'qfrc_smooth')
-  def test_step(self, fname):
+    _assert_attr_eq(d, dx, 'qacc_smooth')
-    """Test mujoco mj step matches mujoco_mjx step."""
-    np.random.seed(test_util.TEST_FILES.index(fname))
-    m = test_util.load_test_file(fname)
-    step_jit_fn = jax.jit(forward.step)
 
-    mx = mjx.device_put(m)
+    # euler
+    dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
+    mujoco.mj_Euler(m, d)
+    _assert_attr_eq(d, dx, 'act')
+    _assert_attr_eq(d, dx, 'qpos')
+    _assert_attr_eq(d, dx, 'time')
+
+  def test_step(self):
+    m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)
-    # give the system a little kick to ensure we have non-identity rotations
+    # apply some control and xfrc input
-    d.qvel = np.random.normal(m.nv) * 0.05
+    d.ctrl = np.array([-18, 0.59, 0.47])
-    for i in range(100):
+    d.xfrc_applied[0, 2] = 0.1  # torque
-      # in order to avoid re-jitting, reuse the same mj_data shape
+    d.xfrc_applied[1, 4] = 0.3  # linear force
-      qpos, qvel = d.qpos, d.qvel
+    mujoco.mj_step(m, d, 100)  # get some dynamics going
-      d = mujoco.MjData(m)
-      d.qpos, d.qvel = qpos, qvel
-      dx = mjx.device_put(d)
 
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
-      dx = step_jit_fn(mx, dx)
+    dx = jax.jit(mjx.step)(mx, mjx.put_data(m, d))
-
+    mujoco.mj_step(m, d)
-      _assert_attr_eq(d, dx, 'qvel', i, fname, atol=1e-2)
+    _assert_attr_eq(d, dx, 'act')
-      _assert_attr_eq(d, dx, 'qpos', i, fname, atol=1e-2)
+    _assert_attr_eq(d, dx, 'time')
-      _assert_attr_eq(d, dx, 'act', i, fname)
+    _assert_attr_eq(d, dx, 'qvel')
-      _assert_attr_eq(d, dx, 'time', i, fname)
+    _assert_attr_eq(d, dx, 'qpos')
 
   def test_rk4(self):
     m = mujoco.MjModel.from_xml_string("""
@@ -94,7 +92,6 @@ class ForwardTest(parameterized.TestCase):
             <flag constraint="disable"/>
           </option>
           <worldbody>
-            <light pos="0 0 1"/>
             <geom type="plane" size="1 1 .01" pos="0 0 -1"/>
             <body pos="0.15 0 0">
               <joint type="hinge" axis="0 1 0"/>
@@ -107,39 +104,33 @@ class ForwardTest(parameterized.TestCase):
           </worldbody>
         </mujoco>
         """)
-    step_jit_fn = jax.jit(forward.step)
 
-    mx = mjx.device_put(m)
     d = mujoco.MjData(m)
     # give the system a little kick to ensure we have non-identity rotations
-    d.qvel = np.random.normal(m.nv) * 0.05
+    d.qvel = np.array([0.2, -0.1])
-    for i in range(100):
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
-      # in order to avoid re-jitting, reuse the same mj_data shape
+    mujoco.mj_forward(m, d)
-      qpos, qvel = d.qpos, d.qvel
-      d = mujoco.MjData(m)
-      d.qpos, d.qvel = qpos, qvel
-      dx = mjx.device_put(d)
 
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
-      dx = step_jit_fn(mx, dx)
+    dx = jax.jit(mjx.rungekutta4)(mx, mjx.put_data(m, d))
+    mujoco.mj_RungeKutta(m, d, 4)
 
-      _assert_attr_eq(d, dx, 'qvel', i, 'test_rk4', atol=1e-2)
+    _assert_attr_eq(d, dx, 'qvel')
-      _assert_attr_eq(d, dx, 'qpos', i, 'test_rk4', atol=1e-2)
+    _assert_attr_eq(d, dx, 'qpos')
-      _assert_attr_eq(d, dx, 'act', i, 'test_rk4')
+    _assert_attr_eq(d, dx, 'act')
-      _assert_attr_eq(d, dx, 'time', i, 'test_rk4')
+    _assert_attr_eq(d, dx, 'time')
 
   def test_disable_eulerdamp(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
-    m.opt.disableflags = m.opt.disableflags | DisableBit.EULERDAMP
+    self.assertTrue((m.dof_damping > 0).any())
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EULERDAMP
 
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    d.qvel[:] = 1.0
-    self.assertTrue((mx.dof_damping > 0).any())
+    d.qacc[:] = 1.0
-    dx = mjx.device_put(d)
+    mx = mjx.put_model(m)
-    dx = jax.jit(forward.forward)(mx, dx)
+    dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
 
-    dx = dx.replace(qvel=jp.ones_like(dx.qvel), qacc=jp.ones_like(dx.qacc))
-    dx = jax.jit(forward._euler)(mx, dx)
     np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep)
 
 

mjx/mujoco/mjx/_src/io.py

@@ -77,6 +77,7 @@ def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic:
 
 def put_model(m: mujoco.MjModel, device=None) -> types.Model:
   """Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
+
   if m.ntendon:
     raise NotImplementedError('tendons are not supported')
 
@@ -150,7 +151,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
   d = types.Data(
       solver_niter=jp.array(0, dtype=jp.int32),
       time=jp.array(0.0),
-      qpos=m.qpos0,
+      qpos=jp.array(m.qpos0),
       qvel=zero_nv,
       act=zero_na,
       qacc_warmstart=zero_nv,

mjx/mujoco/mjx/_src/passive.py

@@ -76,7 +76,7 @@ def _inertia_box_fluid_model(
 def passive(m: Model, d: Data) -> Data:
   """Adds all passive forces."""
   if m.opt.disableflags & DisableBit.PASSIVE:
-    return d
+    return d.replace(qfrc_passive=jp.zeros(m.nv))
 
   # joint-level springs
   def fn(jnt_typs, stiffness, qpos_spring, qpos):

mjx/mujoco/mjx/_src/passive_test.py

@@ -14,100 +14,65 @@
 # ==============================================================================
 """Tests passive forces."""
 
-import itertools
-
 from absl.testing import absltest
-from absl.testing import parameterized
-from etils import epath
 import jax
-import jax.numpy as jp
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import test_util
 import numpy as np
 
-
+# tolerance for difference between MuJoCo and MJX passive calculations - mostly
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-4, rtol=1e-4):
+# due to float precision
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+_TOLERANCE = 1e-7
-  a, b = getattr(a, attr), getattr(b, attr)
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-class PassiveTest(parameterized.TestCase):
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  @parameterized.parameters(enumerate(('ant.xml', 'pendula.xml')))
-  def test_stiffness_damping(self, seed, fname):
-    """Tests stiffness and damping on Ant."""
-    np.random.seed(seed)
-    path = epath.resource_path('mujoco.mjx') / 'test_data'
-    path /= fname
-    m = mujoco.MjModel.from_xml_string(path.read_text())
 
-    # set stiffness/damping
+def _assert_attr_eq(a, b, attr):
-    m.jnt_stiffness = np.random.uniform(size=m.njnt)
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
-    m.dof_damping = np.random.uniform(size=m.nv)
+
+
+class PassiveTest(absltest.TestCase):
+
+  def test_passive(self):
+    m = test_util.load_test_file('pendula.xml')
     d = mujoco.MjData(m)
-    d.qvel = np.random.random(m.nv)  # random kick
+    # 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])
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
 
-    mx = mjx.device_put(m)
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-    dx = mjx.make_data(mx)
+    _assert_attr_eq(d, dx, 'qfrc_passive')
 
-    passive_jit_fn = jax.jit(mjx.passive)
+    # test with fluid forces
+    m.opt.density = 0.01
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_passive')
 
-    for i in range(100):
+    m.opt.viscosity = 0.02
-      qpos, qvel = d.qpos.copy(), d.qvel.copy()
+    mujoco.mj_forward(m, d)
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
-      dx = passive_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-      _assert_attr_eq(d, dx, 'qfrc_passive', i, fname)
+    _assert_attr_eq(d, dx, 'qfrc_passive')
 
-  @parameterized.parameters(
+    m.opt.wind = np.array([0.03, 0.04, 0.05])
-      itertools.product(range(3), ('pendula.xml',))
+    mujoco.mj_forward(m, d)
-  )
+    mx = mjx.put_model(m)
-  def test_fluid(self, seed, fname):
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-    np.random.seed(seed)
+    _assert_attr_eq(d, dx, 'qfrc_passive')
-    path = epath.resource_path('mujoco.mjx') / 'test_data'
-    path /= fname
-    m = mujoco.MjModel.from_xml_string(path.read_text())
 
-    # set density/viscosity/wind
+    # test disable passive
-    m.opt.density = np.random.uniform()
+    mx = mx.tree_replace({'opt.disableflags': mjx.DisableBit.PASSIVE})
-    m.opt.viscosity = np.random.uniform()
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-    m.opt.wind = np.random.uniform()
+    np.testing.assert_allclose(dx.qfrc_passive, 0)
-
-    passive_jit_fn = jax.jit(mjx.passive)
-
-    mx = mjx.device_put(m)
-    d = mujoco.MjData(m)
-    d.qvel = np.random.random(m.nv)  # random kick
-
-    for i in range(100):
-      mujoco.mj_step(m, d)
-      dx = mjx.device_put(d)
-      mujoco.mj_passive(m, d)
-      dx = passive_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'qfrc_passive', i, fname)
-
-  def test_disable_passive(self):
-    m = mujoco.MjModel.from_xml_string("""
-        <mujoco>
-          <option density="1" viscosity="2" wind="0.1 0.2 0.3">
-            <flag passive="disable"/>
-          </option>
-          <worldbody>
-            <body>
-              <joint damping="1" axis="1 0 0" type="ball"/>
-              <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-            </body>
-          </worldbody>
-        </mujoco>
-        """)
-    mx = mjx.device_put(m)
-    d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
-    dx = dx.replace(qvel=jp.ones(mx.nv))
-
-    passive_jit_fn = jax.jit(mjx.passive)
-    dx = passive_jit_fn(mx, dx)
-    np.testing.assert_equal(dx.qfrc_passive, np.zeros(mx.nv))
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/scan_test.py

@@ -193,7 +193,7 @@ class ScanTest(absltest.TestCase):
     </mujoco>
   """
 
-  def testscan_actuators(self):
+  def test_scan_actuators(self):
     """Tests scanning over actuators."""
     m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML)
     m = mjx.device_put(m)

mjx/mujoco/mjx/_src/smooth.py

@@ -435,40 +435,38 @@ def transmission(m: Model, d: Data) -> Data:
   if not m.nu:
     return d
 
-  def fn(gear, jnt_typ, m_i, m_j, qpos):
+  def fn(gear, jnt_typ, m_j, qpos):
     # handles joint transmissions only
     if jnt_typ == JointType.FREE:
       length = jp.zeros(1)
       moment = gear
-      m_i = jp.repeat(m_i, 6)
       m_j = m_j + jp.arange(6)
     elif jnt_typ == JointType.BALL:
-      axis, _ = math.quat_to_axis_angle(qpos)
+      axis, angle = math.quat_to_axis_angle(qpos)
-      length = jp.dot(axis, gear[:3])[None]
+      length = jp.dot(axis * angle, gear[:3])[None]
       moment = gear[:3]
-      m_i = jp.repeat(m_i, 3)
       m_j = m_j + jp.arange(3)
     elif jnt_typ in (JointType.SLIDE, JointType.HINGE):
       length = qpos * gear[0]
       moment = gear[:1]
-      m_i, m_j = m_i[None], m_j[None]
+      m_j = m_j[None]
     else:
       raise RuntimeError(f'unrecognized joint type: {jnt_typ}')
-    return length, moment, m_i, m_j
+    moment = jp.zeros((m.nv,)).at[m_j].set(moment)
+    return length, moment
 
-  length, m_val, m_i, m_j = scan.flat(
+  length, moment = scan.flat(
       m,
       fn,
-      'ujujq',
+      'ujjq',
-      'uvvv',
+      'uuuu',
       m.actuator_gear,
       m.jnt_type,
-      jp.arange(m.nu),
       jp.array(m.jnt_dofadr),
       d.qpos,
       group_by='u',
   )
-  moment = jp.zeros((m.nu, m.nv)).at[m_i, m_j].set(m_val)
   length = length.reshape((m.nu,))
+  moment = moment.reshape((m.nu, m.nv))
   d = d.replace(actuator_length=length, actuator_moment=moment)
   return d

mjx/mujoco/mjx/_src/smooth_test.py

@@ -15,122 +15,107 @@
 """Tests for smooth dynamics functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
 import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.types import DisableBit
-# pylint: enable=g-importing-member
 import numpy as np
 
-
+# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
-def _assert_eq(a, b, name, step, fname, atol=5e-4, rtol=5e-4):
+# due to float precision
-  err_msg = f'mismatch: {name} at step {step} in {fname}'
+_TOLERANCE = 5e-5
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=5e-4, rtol=5e-4):
+def _assert_eq(a, b, name):
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+  tol = _TOLERANCE * 10  # avoid test noise
-  a, b = getattr(a, attr), getattr(b, attr)
+  err_msg = f'mismatch: {name}'
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
 
-class SmoothTest(parameterized.TestCase):
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
 
-  @parameterized.parameters(enumerate(test_util.TEST_FILES))
-  def test_smooth(self, seed, fname):
-    """Tests mujoco mj smooth functions match mujoco_mjx smooth functions."""
-    if fname in ('convex.xml', 'equality.xml'):
-      return
 
-    np.random.seed(seed)
+class SmoothTest(absltest.TestCase):
 
-    m = test_util.load_test_file(fname)
+  def setUp(self):
+    super().setUp()
+    # although we already have generous padding of thresholds, it doesn't hurt
+    # to also fix the seed to reduce test flakiness
+    np.random.seed(0)
+
+  def test_smooth(self):
+    """Tests MJX smooth functions match MuJoCo smooth functions."""
+
+    m = test_util.load_test_file('pendula.xml')
     d = mujoco.MjData(m)
-
-    kinematics_jit_fn = jax.jit(mjx.kinematics)
-    com_pos_jit_fn = jax.jit(mjx.com_pos)
-    crb_jit_fn = jax.jit(mjx.crb)
-    factor_m_fn = jax.jit(mjx.factor_m)
-    com_vel_jit_fn = jax.jit(mjx.com_vel)
-    rne_jit_fn = jax.jit(mjx.rne)
-    mul_m_jit_fn = jax.jit(mjx.mul_m)
-    transmission_jit_fn = jax.jit(mjx.transmission)
-
-    mx = mjx.device_put(m)
-    dx = mjx.make_data(mx)
-
     # give the system a little kick to ensure we have non-identity rotations
     d.qvel = np.random.random(m.nv)
-    for i in range(100):
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
-      qpos, qvel = d.qpos.copy(), d.qvel.copy()
+    mujoco.mj_forward(m, d)
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
 
-      # kinematics
+    # kinematics
-      dx = kinematics_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
+    dx = jax.jit(mjx.kinematics)(mx, mjx.put_data(m, d))
-      _assert_attr_eq(d, dx, 'xanchor', i, fname)
+    _assert_attr_eq(d, dx, 'xanchor')
-      _assert_attr_eq(d, dx, 'xaxis', i, fname)
+    _assert_attr_eq(d, dx, 'xaxis')
-      _assert_attr_eq(d, dx, 'xpos', i, fname)
+    _assert_attr_eq(d, dx, 'xpos')
-      _assert_attr_eq(d, dx, 'xquat', i, fname)
+    _assert_attr_eq(d, dx, 'xquat')
-      _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat', i, fname)
+    _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat')
-      _assert_attr_eq(d, dx, 'xipos', i, fname)
+    _assert_attr_eq(d, dx, 'xipos')
-      _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat', i, fname)
+    _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat')
-      _assert_attr_eq(d, dx, 'geom_xpos', i, fname)
+    _assert_attr_eq(d, dx, 'geom_xpos')
-      _assert_eq(
+    _assert_eq(d.geom_xmat.reshape((-1, 3, 3)), dx.geom_xmat, 'geom_xmat')
-          d.geom_xmat.reshape((-1, 3, 3)),
+    _assert_attr_eq(d, dx, 'site_xpos')
-          dx.geom_xmat,
+    _assert_eq(d.site_xmat.reshape((-1, 3, 3)), dx.site_xmat, 'site_xmat')
-          'geom_xmat',
+    # com_pos
-          i,
+    dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
-          fname,
+    _assert_attr_eq(d, dx, 'subtree_com')
-      )
+    _assert_attr_eq(d, dx, 'cinert')
+    _assert_attr_eq(d, dx, 'cdof')
+    # crb
+    dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'crb')
+    _assert_attr_eq(d, dx, 'qM')
+    # factor_m
+    dx = mjx.put_data(m, d)
+    dx = jax.jit(mjx.factor_m)(mx, dx, dx.qM)
+    _assert_attr_eq(d, dx, 'qLD')
+    _assert_attr_eq(d, dx, 'qLDiagInv')
+    # com_vel
+    dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'cvel')
+    _assert_attr_eq(d, dx, 'cdof_dot')
+    # rne
+    dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_bias')
+    # transmission
+    dx = jax.jit(mjx.transmission)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx, 'actuator_moment')
 
-      # com_pos
+  def test_mul_m(self):
-      dx = com_pos_jit_fn(mx, dx)
+    m = test_util.load_test_file('pendula.xml')
-      _assert_attr_eq(d, dx, 'subtree_com', i, fname)
+    d = mujoco.MjData(m)
-      _assert_attr_eq(d, dx, 'cinert', i, fname)
+    # give the system a little kick to ensure we have non-identity rotations
-      _assert_attr_eq(d, dx, 'cdof', i, fname)
+    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)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    vec = np.random.random(m.nv)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    mj_vec = np.zeros(m.nv)
+    mujoco.mj_mulM(m, d, mj_vec, vec)
+    _assert_eq(mj_vec, mjx_vec, 'mul_m')
 
-      # crb
+  def test_disable_gravity(self):
-      dx = crb_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'crb', i, fname)
-      _assert_attr_eq(d, dx, 'qM', i, fname)
-
-      # factor_m
-      dx = factor_m_fn(mx, dx, dx.qM)
-      _assert_attr_eq(d, dx, 'qLD', i, fname, atol=1e-3)
-      _assert_attr_eq(d, dx, 'qLDiagInv', i, fname, atol=1e-3)
-
-      # com_vel
-      dx = com_vel_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'cvel', i, fname)
-      _assert_attr_eq(d, dx, 'cdof_dot', i, fname)
-
-      # rne
-      dx = rne_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'qfrc_bias', i, fname)
-
-      # mul_m (auxilliary function, not part of smooth step)
-      vec = np.random.random(m.nv)
-      mjx_vec = mul_m_jit_fn(mx, dx, jp.array(vec))
-      mj_vec = np.zeros(m.nv)
-      mujoco.mj_mulM(m, d, mj_vec, vec)
-      _assert_eq(mj_vec, mjx_vec, 'mul_m', i, fname)
-
-      # transmission
-      dx = transmission_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'actuator_length', i, fname)
-      _assert_attr_eq(d, dx, 'actuator_moment', i, fname)
-
-
-class DisableGravityTest(absltest.TestCase):
-
-  def test_disabled(self):
     m = mujoco.MjModel.from_xml_string("""
         <mujoco>
-          <option timestep="0.01"/>
+          <option>
+            <flag gravity="disable"/>
+          </option>
           <worldbody>
             <body>
               <joint type="free"/>
@@ -139,63 +124,13 @@ class DisableGravityTest(absltest.TestCase):
           </worldbody>
         </mujoco>
         """)
-    mx = mjx.device_put(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
-
-    # test with gravity
-    step_jit_fn = jax.jit(mjx.step)
-    dx = step_jit_fn(mx, dx)
-    np.testing.assert_array_almost_equal(
-        dx.qpos, np.array([0.0, 0.0, -9.81e-4, 1.0, 0.0, 0.0, 0.0]), decimal=7
-    )
-
-    # test with gravity disabled
-    mx = mx.tree_replace(
-        {'opt.disableflags': mx.opt.disableflags | DisableBit.GRAVITY}
-    )
-    dx = mjx.device_put(d)
-    step_jit_fn = jax.jit(mjx.step)
-    dx = step_jit_fn(mx, dx)
-    np.testing.assert_equal(
-        dx.qpos, np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0])
-    )
-
-
-class SiteTest(absltest.TestCase):
-
-  def test_site(self):
-    """Tests that site positions and orientations match MuJoCo."""
-    m = mujoco.MjModel.from_xml_string("""
-      <mujoco>
-        <worldbody>
-          <site name="origin"/>
-          <body>
-            <joint type="free"/>
-            <geom pos="1 0 1" type="box" size="0.1 0.01 0.01"/>
-            <site name="s1" pos="1.3 0 0"/>     <!-- pos only -->
-            <site name="s2"/>                   <!-- no pos, no quat -->
-            <site name="s3" quat="1 0 1 0"/>    <!- quat only -->
-            <site name="s4" pos="0 1.5 0" quat="1 0 1 0"/>
-            <site name="s5" pos="1 0 1"/>       <!-- same as ipos -->
-            <site/>
-          </body>
-        </worldbody>
-      </mujoco>
-      """)
-    d = mujoco.MjData(m)
-
-    mx = mjx.device_put(m)
-    dx = mjx.device_put(d)
-
     mujoco.mj_forward(m, d)
-    dx = mjx.forward(mx, dx)
+    mx = mjx.put_model(m)
-
+    dx = mjx.put_data(m, d)
-    np.testing.assert_array_almost_equal(dx.site_xpos, d.site_xpos)
-    np.testing.assert_array_almost_equal(
-        dx.site_xmat, d.site_xmat.reshape((-1, 3, 3))
-    )
 
+    dx = jax.jit(mjx.rne)(mx, dx)
+    np.testing.assert_allclose(dx.qfrc_bias, 0)
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/solver_test.py

@@ -12,119 +12,64 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for forward functions."""
+"""Tests for constraint functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
-from etils import epath
 import jax
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import test_util
 import numpy as np
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-2, rtol=1e-2):
+# tolerance for difference between MuJoCo and MJX constraint calculations,
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+# mostly due to float precision
-  a, b = getattr(a, attr), getattr(b, attr)
+_TOLERANCE = 5e-5
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-class Solver64Test(parameterized.TestCase):
+def _assert_eq(a, b, name, tol=_TOLERANCE):
-  """Tests solvers at 64 bit precision."""
+  tol = tol * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  def setUp(self):
-    super().setUp()
-    jax.config.update('jax_enable_x64', True)
 
-  def tearDown(self):
+def _assert_attr_eq(a, b, attr):
-    super().tearDown()
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
-    jax.config.update('jax_enable_x64', False)
 
-  @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml')))
+
-  def test_cg(self, seed, fname):
+class SolverTest(absltest.TestCase):
-    """Test mjx cg solver matches mujoco cg solver at 64 bit precision."""
+
-    f = epath.resource_path('mujoco.mjx') / 'test_data' / fname
+  def test_solver(self):
-    m = mujoco.MjModel.from_xml_string(f.read_text())
+    """Test solver."""
+    m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
 
-    jax.config.update('jax_enable_x64', True)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
-    forward_jit_fn = jax.jit(mjx.forward)
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
+    _assert_attr_eq(d, dx, 'qacc')
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
+    nnz = dx.efc_J.any(axis=1)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
 
-    # give the system a little kick to ensure we have non-identity rotations
+    # also test normal CG
-    np.random.seed(seed)
+    m.opt.solver = mujoco.mjtSolver.mjSOL_CG
-    d.qvel = 0.01 * np.random.random(m.nv)
+    mujoco.mj_forward(m, d)
-
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
-    for i in range(100):
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
-      # in order to avoid re-jitting, reuse the same mj_data shape
+    _assert_attr_eq(d, dx, 'qacc')
-      save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
-      d = mujoco.MjData(m)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
-      d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save
-      dx = mjx.device_put(d)
-
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx)
-
-      # at 64 bits the solutions returned by the two solvers are quite close
-      self.assertLessEqual(dx.solver_niter[0], d.solver_niter[0])
-      _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname)
-      _assert_attr_eq(d, dx, 'qacc', i, fname)
-
-
-class SolverTest(parameterized.TestCase):
-
-  @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml')))
-  def test_cg(self, seed, fname):
-    """Test mjx cg solver is close to mj at 32 bit precision.
-
-    Args:
-      seed: int
-      fname: file to test
-
-    At lower float resolution there's wiggle room in valid forces that satisfy
-    constraints.  So instead let's mainly validate that mjx is finding solutions
-    with as good cost as mujoco, even if the resulting forces/accelerations
-    are not quite the same.
-    """
-    f = epath.resource_path('mujoco.mjx') / 'test_data' / fname
-    m = mujoco.MjModel.from_xml_string(f.read_text())
-    d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
-
-    forward_jit_fn = jax.jit(mjx.forward)
-
-    # give the system a little kick to ensure we have non-identity rotations
-    np.random.seed(seed)
-    d.qvel = 0.01 * np.random.random(m.nv)
-
-    for i in range(100):
-      # in order to avoid re-jitting, reuse the same mj_data shape
-      save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth
-      d = mujoco.MjData(m)
-      d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save
-      dx = mjx.device_put(d)
-
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx)
-
-      def cost(qacc):
-        jaref = np.zeros(d.nefc)
-        mujoco.mj_mulJacVec(m, d, jaref, qacc)
-        jaref -= d.efc_aref
-        cost = np.array([0.0])
-        mujoco.mj_constraintUpdate(m, d, jaref, cost, 0)
-        return cost[0]
-
-      cost_mj, cost_mjx = cost(d.qacc), cost(dx.qacc)
-
-      self.assertLessEqual(
-          cost_mjx,
-          cost_mj * 1.01,
-          msg=f'mismatch: {fname} at step {i}, cost too high',
-      )
-      _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname, atol=1e-1, rtol=1e-1)
-      _assert_attr_eq(d, dx, 'qacc', i, fname, atol=1e-1, rtol=1e-1)
 
+    # without warmstart, the solution is not as close
+    m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
+    m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-2)
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/support_test.py

@@ -34,8 +34,8 @@ class SupportTest(parameterized.TestCase):
     m = test_util.load_test_file(fname)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
     point = np.random.randn(3)
     body = np.random.choice(m.nbody)
     jacp, jacr = jax.jit(support.jac)(mx, dx, point, body)
@@ -49,11 +49,11 @@ class SupportTest(parameterized.TestCase):
     """Tests that xfrc_accumulate ouput matches mj_xfrcAccumulate."""
     np.random.seed(0)
 
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
     self.assertFalse((dx.xipos == 0.0).all())
 
     xfrc = np.random.rand(*dx.xfrc_applied.shape)

mjx/mujoco/mjx/_src/test_util.py

@@ -23,10 +23,8 @@ import mujoco
 import numpy as np
 
 TEST_FILES: List[str] = [
-    'ant.xml',
+    'constraints.xml',
     'convex.xml',
-    'equality.xml',
-    'humanoid.xml',
     'pendula.xml',
 ]
 

mjx/mujoco/mjx/integration_test/collision_driver_test.py

@@ -58,9 +58,9 @@ class CollisionDriverIntegrationTest(parameterized.TestCase):
     )
 
     m = mujoco.MjModel.from_xml_string(mjcf)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_step(m, d)
     collision_jit_fn = jax.jit(mjx.collision)

mjx/mujoco/mjx/integration_test/forward_test.py

@@ -19,7 +19,6 @@ from absl.testing import parameterized
 import jax
 import mujoco
 from mujoco import mjx
-from mujoco.mjx._src import forward
 from mujoco.mjx._src import test_util
 import numpy as np
 
@@ -46,7 +45,7 @@ class ActuationIntegrationTest(parameterized.TestCase):
         enable_contact=False,
     )
     m = mujoco.MjModel.from_xml_string(mjcf)
-    actuation_jit_fn = jax.jit(forward._actuation)
+    actuation_jit_fn = jax.jit(mjx.fwd_actuation)
 
     # init
     d = mujoco.MjData(m)
@@ -57,8 +56,8 @@ class ActuationIntegrationTest(parameterized.TestCase):
     mujoco.mj_fwdVelocity(m, d)
 
     # put on device
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_fwdActuation(m, d)
     dx = actuation_jit_fn(mx, dx)

mjx/mujoco/mjx/integration_test/smooth_test.py

@@ -60,8 +60,8 @@ class TransmissionIntegrationTest(parameterized.TestCase):
     d.qvel = np.random.random(m.nv)
 
     # put on device
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_transmission(m, d)
     dx = transmission_jit_fn(mx, dx)

mjx/mujoco/mjx/test_data/ant.xml

@@ -1,82 +0,0 @@
-<mujoco model="ant">
-  <compiler angle="degree" coordinate="local" inertiafromgeom="true"/>
-  <option timestep="0.01" iterations="4" solver="CG" />
-  <default>
-    <joint armature="1" damping="1" limited="true"/>
-    <geom contype="0" conaffinity="0" condim="3" density="5.0" friction="1 0.5 0.5"/>
-  </default>
-  <asset>
-    <texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100"/>
-    <texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127"/>
-    <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
-    <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
-    <material name="geom" texture="texgeom" texuniform="true"/>
-  </asset>
-  <worldbody>
-    <light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
-    <geom conaffinity="1" condim="3" material="MatPlane" name="floor" pos="0 0 0" size="40 40 40" type="plane"/>
-    <body name="torso" pos="0 0 0.75">
-      <camera name="track" mode="trackcom" pos="0 -3 0.3" xyaxes="1 0 0 0 0 1"/>
-      <geom name="torso_geom" pos="0 0 0" size="0.25" type="sphere"/>
-      <joint armature="0" damping="0" limited="false" margin="0.01" name="root" pos="0 0 0" type="free"/>
-      <body name="front_left_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="aux_1_geom" size="0.08" type="capsule"/>
-        <body name="aux_1" pos="0.2 0.2 0">
-          <joint axis="0 0 1" name="hip_1" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="left_leg_geom" size="0.08" type="capsule"/>
-          <body pos="0.2 0.2 0">
-            <joint axis="-1 1 0" name="ankle_1" pos="0.0 0.0 0.0" range="30 70" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 0.4 0.4 0.0" name="left_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="left_foot_geom" contype="1" pos="0.4 0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-      <body name="front_right_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="aux_2_geom" size="0.08" type="capsule"/>
-        <body name="aux_2" pos="-0.2 0.2 0">
-          <joint axis="0 0 1" name="hip_2" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="right_leg_geom" size="0.08" type="capsule"/>
-          <body pos="-0.2 0.2 0">
-            <joint axis="1 1 0" name="ankle_2" pos="0.0 0.0 0.0" range="-70 -30" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 -0.4 0.4 0.0" name="right_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="right_foot_geom" contype="1" pos="-0.4 0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-      <body name="back_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="aux_3_geom" size="0.08" type="capsule"/>
-        <body name="aux_3" pos="-0.2 -0.2 0">
-          <joint axis="0 0 1" name="hip_3" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="back_leg_geom" size="0.08" type="capsule"/>
-          <body pos="-0.2 -0.2 0">
-            <joint axis="-1 1 0" name="ankle_3" pos="0.0 0.0 0.0" range="-70 -30" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 -0.4 -0.4 0.0" name="third_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="third_foot_geom" contype="1" pos="-0.4 -0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-      <body name="right_back_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="aux_4_geom" size="0.08" type="capsule"/>
-        <body name="aux_4" pos="0.2 -0.2 0">
-          <joint axis="0 0 1" name="hip_4" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="rightback_leg_geom" size="0.08" type="capsule"/>
-          <body pos="0.2 -0.2 0">
-            <joint axis="1 1 0" name="ankle_4" pos="0.0 0.0 0.0" range="30 70" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 0.4 -0.4 0.0" name="fourth_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="fourth_foot_geom" contype="1" pos="0.4 -0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-  <actuator>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_4" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_4" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_1" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_1" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_2" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_2" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_3" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_3" gear="150"/>
-  </actuator>
-</mujoco>

mjx/mujoco/mjx/test_data/constraints.xml

@@ -0,0 +1,53 @@
+<!-- For validating constraint dynamics:
+
+* connect, weld, joint constraints
+* collision and joint limits for ball and 1d joints
+* solref, solimp
+-->
+<mujoco>
+  <option timestep="0.02"/>
+
+  <default>
+    <default class="box">
+      <geom type="box" size=".2" fromto="0 0 0 0 -2 0" rgba=".4 .7 .6 .3"  contype="0"/>
+    </default>
+  </default>
+
+  <worldbody>
+    <geom pos="0 0 -1" type="plane" size="10 10 .01"/>
+
+    <body name="anchor1" pos="-3 0 0"/>
+    <body name="beam1" pos="-3 0 0">
+      <joint name="joint1" type="ball" range="0 45" solreflimit="0.03 0.9" solimplimit="0.89 0.9 0.01 2.1"/>
+      <geom class="box"/>
+    </body>
+
+    <body name="anchor2" pos="-1 0 0"/>
+    <body name="beam2" pos="-1 0 0">
+      <freejoint/>
+      <geom class="box"/>
+    </body>
+
+    <body name="beam3" pos="1 0 0">
+      <joint name="joint3" axis="1 0 0" type="hinge" range="-20 20"/>
+      <geom class="box"/>
+    </body>
+
+    <body name="beam4" pos="3 0 0">
+      <joint name="joint4" axis="1 0 0" type="hinge" damping="1.0"/> <!-- tests no joint range -->
+      <geom class="box"/>
+    </body>
+  </worldbody>
+
+  <equality>
+    <connect name="connect"  body1="anchor1"    body2="beam1"   anchor="1 0 -1" />
+    <weld    name="weld"     body1="anchor2"    body2="beam2"   relpose="0 0 0 1 -.3 0 0" torquescale="0.002" anchor="0 -2 0"/>
+    <joint   name="joint"    joint1="joint3"    joint2="joint4" polycoef="0.5 -1 0.1 0.15 0.2" />
+  </equality>
+
+  <actuator>
+    <position ctrlrange="-20 20" gear="500" joint="joint1" name="act1"/>
+    <motor gear="50000" joint="joint3" name="act2"/>
+    <motor gear="75000" joint="joint4" name="act3"/>
+  </actuator>
+</mujoco>

mjx/mujoco/mjx/test_data/equality.xml

@@ -1,71 +0,0 @@
-<mujoco>
-  <option solver="CG" iterations="8">
-    <flag contact="disable"/>
-  </option>
-
-  <default>
-    <default class="static">
-      <geom size=".5 .1 .5" rgba=".5 .7 .5 .3"/>
-    </default>
-    <default class="free">
-      <geom type="box" size=".2" fromto="0 0 0 0 -2 0" rgba=".4 .7 .6 .3"/>
-    </default>
-  </default>
-
-  <worldbody>
-    <geom pos="0 0 -2" type="plane" size="10 10 .01"/>
-    <light pos="0 0 20"/>
-
-    <body name="box1" pos="-3 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam1" pos="-3 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box2" pos="-1 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam2" pos="-1 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box3" pos="1 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam3" pos="1 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box4" pos="3 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam4" pos="3 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box5" pos="4 0 0">
-      <geom class="free"/>
-      <joint name="joint1" axis="1 0 0" type="hinge" />
-    </body>
-
-    <body name="box6" pos="4 0 0">
-      <geom class="free"/>
-      <joint name="joint2" axis="1 0 0" type="hinge" />
-    </body>
-
-
-  </worldbody>
-
-  <equality>
-    <connect name="connect anchor"           body1="box1"    body2="beam1"   anchor="0 0 -1" />
-    <weld    name="weld anchor weak torques" body1="box2"    body2="beam2"   torquescale="0.002" anchor="0 -2 0"/>
-    <weld    name="weld relpose"             body1="box3"    body2="beam3"   relpose="0 0 0 1 -.3 0 0"/>
-    <weld    name="weld relpose+anchor"      body1="box4"    body2="beam4"   relpose="0 0 0 1 -.3 0 0" anchor="0 0 -1"/>
-    <joint   name="joint"                    joint1="joint1" joint2="joint2" polycoef="0 -1 0.1 0.15 0.2" />
-  </equality>
-</mujoco>

mjx/mujoco/mjx/test_data/humanoid.xml

@@ -1,109 +0,0 @@
-<mujoco model="humanoid">
-    <compiler angle="degree" inertiafromgeom="true"/>
-    <default>
-        <joint armature="1" damping="1" limited="true"/>
-        <geom conaffinity="0" condim="3" contype="0" material="geom"/>
-        <motor ctrllimited="true" ctrlrange="-.4 .4"/>
-    </default>
-    <option iterations="8" timestep="0.003"/>
-    <size nkey="5" nuser_geom="1"/>
-    <visual>
-        <map fogend="5" fogstart="3"/>
-    </visual>
-    <asset>
-        <texture builtin="gradient" height="100" rgb1=".4 .5 .6" rgb2="0 0 0" type="skybox" width="100"/>
-        <!-- <texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100"/>-->
-        <texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127"/>
-        <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
-        <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
-        <material name="geom" texture="texgeom" texuniform="true"/>
-    </asset>
-    <worldbody>
-        <light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
-        <geom conaffinity="1" condim="3" friction="1 .1 .1" material="MatPlane" name="floor" pos="0 0 0" size="20 20 0.125" type="plane"/>
-        <!-- <geom condim="3" material="MatPlane" name="floor" pos="0 0 0" size="10 10 0.125" type="plane"/>-->
-        <body name="torso" pos="0 0 1.4">
-            <camera name="track" mode="trackcom" pos="0 -4 0" xyaxes="1 0 0 0 0 1"/>
-            <joint armature="0" damping="0" limited="false" name="root" pos="0 0 0" stiffness="0" type="free"/>
-            <geom fromto="0 -.07 0 0 .07 0" name="torso1" size="0.07" type="capsule"/>
-            <geom name="head" pos="0 0 .19" size=".09" type="sphere" user="258"/>
-            <geom fromto="-.01 -.06 -.12 -.01 .06 -.12" name="uwaist" size="0.06" type="capsule"/>
-            <body name="lwaist" pos="-.01 0 -0.260" quat="1.000 0 -0.002 0">
-                <geom fromto="0 -.06 0 0 .06 0" name="lwaist" size="0.06" type="capsule"/>
-                <joint armature="0.02" axis="0 0 1" damping="5" name="abdomen_z" pos="0 0 0.065" range="-45 45" stiffness="20" type="hinge"/>
-                <joint armature="0.02" axis="0 1 0" damping="5" name="abdomen_y" pos="0 0 0.065" range="-75 30" stiffness="10" type="hinge"/>
-                <body name="pelvis" pos="0 0 -0.165" quat="1.000 0 -0.002 0">
-                    <joint armature="0.02" axis="1 0 0" damping="5" name="abdomen_x" pos="0 0 0.1" range="-35 35" stiffness="10" type="hinge"/>
-                    <geom fromto="-.02 -.07 0 -.02 .07 0" name="butt" size="0.09" type="capsule"/>
-                    <body name="right_thigh" pos="0 -0.1 -0.04">
-                        <joint armature="0.01" axis="1 0 0" damping="5" name="right_hip_x" pos="0 0 0" range="-25 5" stiffness="10" type="hinge"/>
-                        <joint armature="0.01" axis="0 0 1" damping="5" name="right_hip_z" pos="0 0 0" range="-60 35" stiffness="10" type="hinge"/>
-                        <joint armature="0.0080" axis="0 1 0" damping="5" name="right_hip_y" pos="0 0 0" range="-110 20" stiffness="20" type="hinge"/>
-                        <geom fromto="0 0 0 0 0.01 -.34" name="right_thigh1" size="0.06" type="capsule"/>
-                        <body name="right_shin" pos="0 0.01 -0.403">
-                            <joint armature="0.0060" axis="0 -1 0" name="right_knee" pos="0 0 .02" range="-160 -2" type="hinge"/>
-                            <geom fromto="0 0 0 0 0 -.3" name="right_shin1" size="0.049" type="capsule"/>
-                            <body name="right_foot" pos="0 0 -0.45">
-                                <geom contype="1" name="right_foot" pos="0 0 0.1" size="0.075" type="sphere" user="0"/>
-                            </body>
-                        </body>
-                    </body>
-                    <body name="left_thigh" pos="0 0.1 -0.04">
-                        <joint armature="0.01" axis="-1 0 0" damping="5" name="left_hip_x" pos="0 0 0" range="-25 5" stiffness="10" type="hinge"/>
-                        <joint armature="0.01" axis="0 0 -1" damping="5" name="left_hip_z" pos="0 0 0" range="-60 35" stiffness="10" type="hinge"/>
-                        <joint armature="0.01" axis="0 1 0" damping="5" name="left_hip_y" pos="0 0 0" range="-110 20" stiffness="20" type="hinge"/>
-                        <geom fromto="0 0 0 0 -0.01 -.34" name="left_thigh1" size="0.06" type="capsule"/>
-                        <body name="left_shin" pos="0 -0.01 -0.403">
-                            <joint armature="0.0060" axis="0 -1 0" name="left_knee" pos="0 0 .02" range="-160 -2" stiffness="1" type="hinge"/>
-                            <geom fromto="0 0 0 0 0 -.3" name="left_shin1" size="0.049" type="capsule"/>
-                            <body name="left_foot" pos="0 0 -0.45">
-                                <geom contype="1" name="left_foot" type="sphere" size="0.075" pos="0 0 0.1" user="0" />
-                            </body>
-                        </body>
-                    </body>
-                </body>
-            </body>
-            <body name="right_upper_arm" pos="0 -0.17 0.06">
-                <joint armature="0.0068" axis="2 1 1" name="right_shoulder1" pos="0 0 0" range="-85 60" stiffness="1" type="hinge"/>
-                <joint armature="0.0051" axis="0 -1 1" name="right_shoulder2" pos="0 0 0" range="-85 60" stiffness="1" type="hinge"/>
-                <geom fromto="0 0 0 .16 -.16 -.16" name="right_uarm1" size="0.04 0.16" type="capsule"/>
-                <body name="right_lower_arm" pos=".18 -.18 -.18">
-                    <joint armature="0.0028" axis="0 -1 1" name="right_elbow" pos="0 0 0" range="-90 50" stiffness="0" type="hinge"/>
-                    <geom fromto="0.01 0.01 0.01 .17 .17 .17" name="right_larm" size="0.031" type="capsule"/>
-                    <geom name="right_hand" pos=".18 .18 .18" size="0.04" type="sphere"/>
-                    <camera pos="0 0 0"/>
-                </body>
-            </body>
-            <body name="left_upper_arm" pos="0 0.17 0.06">
-                <joint armature="0.0068" axis="2 -1 1" name="left_shoulder1" pos="0 0 0" range="-60 85" stiffness="1" type="hinge"/>
-                <joint armature="0.0051" axis="0 1 1" name="left_shoulder2" pos="0 0 0" range="-60 85" stiffness="1" type="hinge"/>
-                <geom fromto="0 0 0 .16 .16 -.16" name="left_uarm1" size="0.04 0.16" type="capsule"/>
-                <body name="left_lower_arm" pos=".18 .18 -.18">
-                    <joint armature="0.0028" axis="0 -1 -1" name="left_elbow" pos="0 0 0" range="-90 50" stiffness="0" type="hinge"/>
-                    <geom fromto="0.01 -0.01 0.01 .17 -.17 .17" name="left_larm" size="0.031" type="capsule"/>
-                    <geom name="left_hand" pos=".18 -.18 .18" size="0.04" type="sphere"/>
-                </body>
-            </body>
-        </body>
-    </worldbody>
-
-    <actuator>
-        <motor gear="100" joint="abdomen_y" name="abdomen_y"/>
-        <motor gear="100" joint="abdomen_z" name="abdomen_z"/>
-        <motor gear="100" joint="abdomen_x" name="abdomen_x"/>
-        <motor gear="100" joint="right_hip_x" name="right_hip_x"/>
-        <motor gear="100" joint="right_hip_z" name="right_hip_z"/>
-        <motor gear="300" joint="right_hip_y" name="right_hip_y"/>
-        <motor gear="200" joint="right_knee" name="right_knee"/>
-        <motor gear="100" joint="left_hip_x" name="left_hip_x"/>
-        <motor gear="100" joint="left_hip_z" name="left_hip_z"/>
-        <motor gear="300" joint="left_hip_y" name="left_hip_y"/>
-        <motor gear="200" joint="left_knee" name="left_knee"/>
-        <motor gear="25" joint="right_shoulder1" name="right_shoulder1"/>
-        <motor gear="25" joint="right_shoulder2" name="right_shoulder2"/>
-        <motor gear="25" joint="right_elbow" name="right_elbow"/>
-        <motor gear="25" joint="left_shoulder1" name="left_shoulder1"/>
-        <motor gear="25" joint="left_shoulder2" name="left_shoulder2"/>
-        <motor gear="25" joint="left_elbow" name="left_elbow"/>
-    </actuator>
-</mujoco>

mjx/mujoco/mjx/test_data/pendula.xml

@@ -18,6 +18,8 @@
   </default>
 
   <worldbody>
+    <site name="origin"/>
+
     <!-- a single free body -->
     <body pos="0 0 0">
       <freejoint/>
@@ -26,45 +28,49 @@
 
     <!-- a single ball joint with a limit -->
     <body pos="0.5 0 0">
-      <joint type="ball" range="0 35"/>
+      <joint name="joint1" type="ball" range="0 35"/>
+      <site name="s1" pos="1.3 0 0"/>
       <geom/>
     </body>
 
     <!-- a single slide joint with a limit -->
     <body pos="1.0 0 0">
-      <joint type="slide" axis="0.1 0.2 0.3" range="-1 1"/>
+      <joint name="joint2" type="slide" axis="0.1 0.2 0.3" range="-1 1"/>
+      <site name="s2"/>
       <geom/>
     </body>
 
     <!-- a single hinge joint with a limit -->
     <body pos="1.5 0 0">
-      <joint type="hinge" axis="0.1 0.2 0.3" range="-35 50"/>
+      <joint name="joint3" type="hinge" axis="0.1 0.2 0.3" range="-35 50"/>
+      <site name="s3" quat="1 0 1 0"/>
       <geom/>
     </body>
 
     <!-- stacked joint: hinge + slide -->
     <body pos="2.0 0 0">
-      <joint type="hinge" axis="0.1 0.2 0.3"/>
+      <joint name="joint4" type="hinge" axis="0.1 0.2 0.3"/>
-      <joint type="slide" axis="0.4 0.5 0.6" range="0 1"/>
+      <joint name="joint5" type="slide" axis="0.4 0.5 0.6" range="0 1"/>
       <geom/>
     </body>
 
     <!-- stacked joint: slide + ball -->
     <body pos="2.5 0 0">
-      <joint type="slide" axis="0.4 0.5 0.6" range="0 1"/>
+      <joint name="joint6" type="slide" axis="0.4 0.5 0.6" range="0 1"/>
       <joint type="ball"/>
       <geom/>
     </body>
 
     <!-- triple pendulum of hinges -->
     <body pos="3.0 0 0">
-      <joint axis="0.1 0.2 0.3" type="hinge"/>
+      <joint name="joint7" axis="0.1 0.2 0.3" type="hinge"/>
       <geom/>
       <body pos="0 0 -0.8">
-        <joint axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
+        <joint name="joint8" axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
         <geom/>
         <body pos="0 -0.7 0">
-          <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <joint name="joint9" axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <site name="s4" pos="0 1.5 0" quat="1 0 1 0"/>
           <geom/>
         </body>
       </body>
@@ -72,14 +78,14 @@
 
     <!-- cherry pendulum: two bodies attached to same parent body -->
     <body pos="3.5 0 0">
-      <joint type="ball" damping="0.5" />
+      <joint name="joint10" type="ball" damping="0.5" />
       <geom/>
       <body pos="0 0 -0.8">
-        <joint axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
+        <joint name="joint11" axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
         <geom/>
       </body>
       <body pos="0 -0.7 0">
-        <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+        <joint name="joint12" axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
         <geom/>
       </body>
     </body>
@@ -89,14 +95,21 @@
       <freejoint/>
       <geom/>
       <body pos="0 0 -0.8">
-        <joint axis="0.4 0.5 0.6" type="slide" armature="0.02" range="-0.4 0.6"/>
+        <joint name="joint13" axis="0.4 0.5 0.6" type="slide" armature="0.02" range="-0.4 0.6"/>
         <geom/>
         <body pos="0 -0.7 0">
-          <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <joint name="joint14" axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
           <geom/>
         </body>
       </body>
     </body>
-
   </worldbody>
+
+  <actuator>
+    <motor gear="250 0 0" joint="joint1" name="act1"/>
+    <motor gear="0 275 0" joint="joint1" name="act2"/>
+    <motor gear="0 0 300" joint="joint1" name="act3"/>
+    <motor gear="275" joint="joint2" name="act4"/>
+    <motor gear="275" joint="joint3" name="act5"/>
+  </actuator>
 </mujoco>