Simplify pendula test data. Fix bug in `scan.body_tree` that led to incorrect smooth dynamics for some kinematic tree layouts.

PiperOrigin-RevId: 582125477
Change-Id: I47505ceb4be4c88052bc670f401b8c45ff320bbd

doc/changelog.rst

@@ -59,6 +59,7 @@ MJX
 - Fixed bug where mixed ``jnt_limited`` joints were not being constrained correctly.
 - Made ``device_put`` type validation more verbose (fixes :github:issue:`1113`).
 - Removed empty EFC rows from ``MJX``, for joints with no limits (fixes :github:issue:`1117`).
+- Fixed bug in ``scan.body_tree`` that led to incorrect smooth dynamics for some kinematic tree layouts.
 
 Python bindings
 ^^^^^^^^^^^^^^^

mjx/mujoco/mjx/_src/constraint.py

@@ -321,10 +321,10 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
   if m.opt.disableflags & DisableBit.EQUALITY:
     ne = 0
   else:
-    ne_weld = (m.eq_type == EqType.WELD).sum()
     ne_connect = (m.eq_type == EqType.CONNECT).sum()
+    ne_weld = (m.eq_type == EqType.WELD).sum()
     ne_joint = (m.eq_type == EqType.JOINT).sum()
-    ne = ne_weld * 6 + ne_connect * 3 + ne_joint
+    ne = ne_connect * 3 + ne_weld * 6 + ne_joint
 
   nf = 0
 

mjx/mujoco/mjx/_src/constraint_test.py

@@ -37,7 +37,7 @@ def _assert_eq(a, b, name, step, fname, atol=5e-3, rtol=5e-3):
 class ConstraintTest(parameterized.TestCase):
 
   @parameterized.parameters(enumerate(test_util.TEST_FILES))
-  def testconstraints(self, seed, fname):
+  def test_constraints(self, seed, fname):
     """Test constraints."""
     np.random.seed(seed)
 

mjx/mujoco/mjx/_src/forward_test.py

@@ -14,8 +14,6 @@
 # ==============================================================================
 """Tests for forward functions."""
 
-import itertools
-
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -38,13 +36,12 @@ def _assert_attr_eq(a, b, attr, step, fname, atol=1e-3, rtol=1e-3):
 
 class ForwardTest(parameterized.TestCase):
 
-  @parameterized.parameters(enumerate(test_util.TEST_FILES))
+  @parameterized.parameters(
-  def test_forward(self, seed, fname):
+      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."""
-    if fname in ('equality.xml',):
+    np.random.seed(test_util.TEST_FILES.index(fname))
-      return
-
-    np.random.seed(seed)
 
     m = test_util.load_test_file(fname)
     d = mujoco.MjData(m)
@@ -62,36 +59,20 @@ class ForwardTest(parameterized.TestCase):
       _assert_attr_eq(d, dx, 'qfrc_smooth', i, fname)
       _assert_attr_eq(d, dx, 'qacc_smooth', i, fname)
 
-  @parameterized.parameters(itertools.product(test_util.TEST_FILES, (0, 1)))
+  @parameterized.parameters(
-  def test_step(self, fname, integrator_type):
+      filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES)
+  )
+  def test_step(self, fname):
     """Test mujoco mj step matches mujoco_mjx step."""
-    if fname in (
+    np.random.seed(test_util.TEST_FILES.index(fname))
-        'mixed_joint_pendulum.xml',
-        'ball_pendulum.xml',
-        'convex.xml',
-        'humanoid.xml',
-        'triple_pendulum.xml',  # TODO(b/301485081)
-        'equality.xml',
-    ):
-      # skip models with big constraint violations at step 0 or too slow to run
-      return
-
-    np.random.seed(integrator_type)
     m = test_util.load_test_file(fname)
     step_jit_fn = jax.jit(forward.step)
 
-    m.opt.integrator = integrator_type
-    int_typ = 'euler' if integrator_type == 0 else 'rk4'
-    test_name = f'{fname} - {int_typ}'
-    steps = 100 if int_typ == 'euler' else 30
-    dt = m.opt.timestep
-    m.opt.timestep = dt if int_typ == 'euler' else dt * 3
-
     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
-    for i in range(steps):
+    for i in range(100):
       # in order to avoid re-jitting, reuse the same mj_data shape
       qpos, qvel = d.qpos, d.qvel
       d = mujoco.MjData(m)
@@ -101,10 +82,51 @@ class ForwardTest(parameterized.TestCase):
       mujoco.mj_step(m, d)
       dx = step_jit_fn(mx, dx)
 
-      _assert_attr_eq(d, dx, 'qvel', i, test_name, atol=1e-2)
+      _assert_attr_eq(d, dx, 'qvel', i, fname, atol=1e-2)
-      _assert_attr_eq(d, dx, 'qpos', i, test_name, atol=1e-2)
+      _assert_attr_eq(d, dx, 'qpos', i, fname, atol=1e-2)
-      _assert_attr_eq(d, dx, 'act', i, test_name)
+      _assert_attr_eq(d, dx, 'act', i, fname)
-      _assert_attr_eq(d, dx, 'time', i, test_name)
+      _assert_attr_eq(d, dx, 'time', i, fname)
+
+  def test_rk4(self):
+    m = mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <option integrator="RK4">
+            <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"/>
+              <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+              <body pos="0.1 0 0">
+                <joint type="slide" axis="1 0 0" stiffness="200"/>
+                <geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
+              </body>
+            </body>
+          </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
+    for i in range(100):
+      # in order to avoid re-jitting, reuse the same mj_data shape
+      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)
+      dx = step_jit_fn(mx, dx)
+
+      _assert_attr_eq(d, dx, 'qvel', i, 'test_rk4', atol=1e-2)
+      _assert_attr_eq(d, dx, 'qpos', i, 'test_rk4', atol=1e-2)
+      _assert_attr_eq(d, dx, 'act', i, 'test_rk4')
+      _assert_attr_eq(d, dx, 'time', i, 'test_rk4')
 
   def test_disable_eulerdamp(self):
     m = test_util.load_test_file('ant.xml')

mjx/mujoco/mjx/_src/passive_test.py

@@ -26,7 +26,7 @@ from mujoco import mjx
 import numpy as np
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5):
+def _assert_attr_eq(a, b, attr, step, fname, atol=1e-4, rtol=1e-4):
   err_msg = f'mismatch: {attr} at step {step} in {fname}'
   a, b = getattr(a, attr), getattr(b, attr)
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
@@ -34,7 +34,7 @@ def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5):
 
 class PassiveTest(parameterized.TestCase):
 
-  @parameterized.parameters(enumerate(('ant.xml', 'mixed_joint_pendulum.xml')))
+  @parameterized.parameters(enumerate(('ant.xml', 'pendula.xml')))
   def test_stiffness_damping(self, seed, fname):
     """Tests stiffness and damping on Ant."""
     np.random.seed(seed)
@@ -60,7 +60,7 @@ class PassiveTest(parameterized.TestCase):
       _assert_attr_eq(d, dx, 'qfrc_passive', i, fname)
 
   @parameterized.parameters(
-      itertools.product(range(3), ('triple_pendulum.xml',))
+      itertools.product(range(3), ('pendula.xml',))
   )
   def test_fluid(self, seed, fname):
     np.random.seed(seed)

mjx/mujoco/mjx/_src/scan.py

@@ -162,7 +162,7 @@ def flat(
 ) -> Y:
   r"""Scan a function across bodies or actuators.
 
-  Scan group data according to type and batch shape then calls vmap(f) on it.\
+  Scan group data according to type and batch shape then calls vmap(f) on it.
 
   Args:
     m: an mjx model
@@ -340,48 +340,88 @@ def body_tree(
       IndexError: if function output shape does not match out_types shape
   """
   _check_input(m, args, in_types)
-  depth_fn = lambda i, p=m.body_parentid: int(i > 0) and 1 + depth_fn(p[i])
-  typ_body_id = {
-      'j': m.jnt_bodyid,
-      'v': m.dof_bodyid,
-      'q': _q_bodyid(m),
-  }
-  key_parents = {}
 
-  # build up groupings of bodies and type ids using (level, (jnt_type,)) keys
+  # group together bodies that will be processed together.  grouping key:
-  key_typ_ids, key_body_ids = {}, {}
+  #   1) the tree depth: parent bodies are processed first, so that they are
-  for body_id in np.arange(m.nbody, dtype=np.int32):
+  #      available as carry input to child bodies (or reverse if reverse=True)
-    depth = depth_fn(body_id)
+  #   2) the types of arguments passed to f, both carry and *args:
+  #      * for 'b' arguments, there is no extra grouping
+  #      * for 'j' arguments, we group by joint type
+  #      * for 'q' arguments, we group by q width
+  #      * for 'v' arguments, we group by dof width
+  depths = np.zeros(m.nbody, dtype=np.int32)
 
-    # create grouping key
+  # map key => body id
-    if any(t in 'jqv' for t in in_types + out_types):
+  key_body_ids = {}
-      jnts = np.nonzero(typ_body_id['j'] == body_id)[0]
+  for body_id in range(m.nbody):
-      jnts_p = np.nonzero(typ_body_id['j'] == m.body_parentid[body_id])[0]
+    parent_id = -1
-      key = depth, tuple(m.jnt_type[jnts])
+    if body_id > 0:
-      parent_key = depth - 1, tuple(m.jnt_type[jnts_p])
+      parent_id = m.body_parentid[body_id]
-    else:
+      depths[body_id] = 1 + depths[parent_id]
-      key, parent_key = (depth, ()), (depth - 1, ())
+
+    # create grouping key: depth, carry, args
+    key = (depths[body_id],)
+
+    for i, t in enumerate(out_types + in_types):
+      id_ = parent_id if i < len(out_types) else body_id
+      if t == 'b':
+        continue
+      elif t == 'j':
+        key += (tuple(m.jnt_type[np.nonzero(m.jnt_bodyid == id_)[0]]))
+      elif t == 'v':
+        key += (len(np.nonzero(m.dof_bodyid == id_)[0]),)
+      elif t == 'q':
+        key += (len(np.nonzero(_q_bodyid(m) == id_)[0]),)
 
-    key_parents[key] = parent_key
     body_ids = key_body_ids.get(key, np.array([], dtype=np.int32))
     key_body_ids[key] = np.append(body_ids, body_id)
 
-    # add ids per type
+  # find parent keys of each key.  a key may have multiple parents if the
-    for t in set(in_types + out_types):
+  # carry output keys of distinct parents are the same. e.g.:
-      out = key_typ_ids.setdefault(key, {})
+  # - depth 0 body 1 (slide joint)
-      id_ = body_id if t == 'b' else np.nonzero(typ_body_id[t] == body_id)[0]
+  # -- depth 1 body 1 (hinge joint)
-      id_ = np.expand_dims(id_, axis=0)
+  # - depth 0 body 2 (ball joint)
-      out[t] = np.concatenate((out[t], id_)) if t in out else id_
+  # -- depth 1 body 2 (hinge joint)
+  # given a scan with 'j' in the in_types, we would group depth 0 bodies
+  # separately but we may group depth 1 bodies together
+  key_parents = {}
 
-  key_typ_ids = list(sorted(key_typ_ids.items(), reverse=reverse))
+  for key, body_ids in key_body_ids.items():
+    body_ids = body_ids[body_ids != 0]  # ignore worldbody, has no parent
+    if body_ids.size == 0:
+      continue
+    # find any key which has a body id that is a parent of these body_ids
+    pids = m.body_parentid[body_ids]
+    parents = {k for k, v in key_body_ids.items() if np.isin(v, pids).any()}
+    key_parents[key] = list(sorted(parents))
+
+  # key => take indices
+  key_in_take, key_y_take = {}, {}
+  for key, body_ids in key_body_ids.items():
+    for i, typ in enumerate(in_types + out_types):
+      if typ == 'b':
+        ids = body_ids
+      elif typ == 'j':
+        ids = np.stack([np.nonzero(m.jnt_bodyid == b)[0] for b in body_ids])
+      elif typ == 'v':
+        ids = np.stack([np.nonzero(m.dof_bodyid == b)[0] for b in body_ids])
+      elif typ == 'q':
+        ids = np.stack([np.nonzero(_q_bodyid(m) == b)[0] for b in body_ids])
+      else:
+        raise ValueError(f'Unknown in_type: {typ}')
+      if i < len(in_types):
+        key_in_take.setdefault(key, []).append(ids)
+      else:
+        key_y_take.setdefault(key, []).append(np.hstack(ids))
 
   # use this grouping to take the right data subsets and call vmap(f)
+  keys = sorted(key_body_ids, reverse=reverse)
   key_y = {}
-  for key, typ_ids in key_typ_ids:
+  for key in keys:
     carry = None
 
     if reverse:
-      child_keys = [k for k, v in key_parents.items() if v == key]
+      child_keys = [k for k, v in key_parents.items() if key in v]
 
       for child_key in child_keys:
         y = key_y[child_key]
@@ -394,39 +434,33 @@ def body_tree(
 
         y = jax.tree_map(index_sum, y)
         carry = y if carry is None else jax.tree_map(jp.add, carry, y)
-    else:
+    elif key in key_parents:
-      parent_key = key_parents[key]
+      ys = [key_y[p] for p in key_parents[key]]
-      y = key_y.get(parent_key)
+      y = jax.tree_map(lambda *x: jp.concatenate(x), *ys)
+      body_ids = np.concatenate([key_body_ids[p] for p in key_parents[key]])
+      parent_ids = m.body_parentid[key_body_ids[key]]
+      take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i)
+      carry = jax.tree_map(take_fn, y)
 
-      if y is not None:
+    f_args = [_take(arg, ids) for arg, ids in zip(args, key_in_take[key])]
-        body_ids = key_body_ids[parent_key]
-        parent_ids = m.body_parentid[key_body_ids[key]]
-        take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i)
-        carry = jax.tree_map(take_fn, y)
-
-    f_args = [_take(arg, typ_ids[typ]) for arg, typ in zip(args, in_types)]
     key_y[key] = _nvmap(f, carry, *f_args)
 
   # slice None results from the final output
-  key_typ_ids = [(k, v) for k, v in key_typ_ids if key_y[k] is not None]
+  keys = [k for k in keys if key_y[k] is not None]
 
-  # concatenate back to a single tree and drop the grouping dimension
+  # concatenate ys, drop grouping dimensions, put back in order
-  ys = [key_y[key] for key, _ in key_typ_ids]
+  y = []
-  f_ret_is_seq = isinstance(ys[0], (list, tuple))
+  for i, typ in enumerate(out_types):
-  ys = ys if f_ret_is_seq else [[y] for y in ys]
+    y_typ = [key_y[key] for key in keys]
-  ys = [
+    if len(out_types) > 1:
-      [v if typ == 'b' else jp.concatenate(v) for v, typ in zip(y, out_types)]
+      y_typ = [y_[i] for y_ in y_typ]
-      for y in ys
+    if typ != 'b':
-  ]
+      y_typ = jax.tree_map(jp.concatenate, y_typ)
-  ys = jax.tree_map(lambda *x: jp.concatenate(x), *ys)
+    y_typ = jax.tree_map(lambda *x: jp.concatenate(x), *y_typ)
+    y_take = np.argsort(np.concatenate([key_y_take[key][i] for key in keys]))
+    _check_output(y_typ, y_take, typ, i)
+    y.append(_take(y_typ, y_take))
 
-  # put concatenated results back into body order
+  y = y[0] if len(out_types) == 1 else y
-  reordered_ys = []
-  for i, (y, typ) in enumerate(zip(ys, out_types)):
-    ids = np.concatenate([np.hstack(v[typ]) for _, v in key_typ_ids])
-    take_ids = _index(ids, np.sort(ids))
-    _check_output(y, take_ids, typ, i)
-    reordered_ys.append(_take(y, take_ids))
-  y = reordered_ys if f_ret_is_seq else reordered_ys[0]
 
   return y

mjx/mujoco/mjx/_src/smooth_test.py

@@ -27,12 +27,12 @@ from mujoco.mjx._src.types import DisableBit
 import numpy as np
 
 
-def _assert_eq(a, b, name, step, fname, atol=1e-5, rtol=1e-5):
+def _assert_eq(a, b, name, step, fname, atol=5e-4, rtol=5e-4):
   err_msg = f'mismatch: {name} at step {step} in {fname}'
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5):
+def _assert_attr_eq(a, b, attr, step, fname, atol=5e-4, rtol=5e-4):
   err_msg = f'mismatch: {attr} at step {step} in {fname}'
   a, b = getattr(a, attr), getattr(b, attr)
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
@@ -101,7 +101,7 @@ class SmoothTest(parameterized.TestCase):
       # 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, rtol=1e-4)
+      _assert_attr_eq(d, dx, 'qLDiagInv', i, fname, atol=1e-3)
 
       # com_vel
       dx = com_vel_jit_fn(mx, dx)
@@ -110,14 +110,14 @@ class SmoothTest(parameterized.TestCase):
 
       # rne
       dx = rne_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'qfrc_bias', i, fname, atol=1e-4)
+      _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, atol=1e-4)
+      _assert_eq(mj_vec, mjx_vec, 'mul_m', i, fname)
 
       # transmission
       dx = transmission_jit_fn(mx, dx)

mjx/mujoco/mjx/_src/test_util.py

@@ -24,16 +24,10 @@ import numpy as np
 
 TEST_FILES: List[str] = [
     'ant.xml',
-    'ball_pendulum.xml',
-    'cherry_pendulum.xml',
     'convex.xml',
     'equality.xml',
     'humanoid.xml',
-    'mixed_joint_pendulum.xml',
+    'pendula.xml',
-    'single_pendulum.xml',
-    'slide_pendulum.xml',
-    'triple_pendulum.xml',
-    'triple_pendulum_free.xml',
 ]
 
 _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']

mjx/mujoco/mjx/test_data/ball_pendulum.xml

@@ -1,23 +0,0 @@
-<mujoco model="ball_pendulum">
-  <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
-  <compiler autolimits="true"/>
-  <default>
-    <geom contype="0" conaffinity="0"/>
-    <joint damping="10"/>
-  </default>
-  <option solver="CG"/>
-  <worldbody>
-    <body>
-      <joint axis="1 0 0" type="ball" range="0 10"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-      <body pos="0.3 0.4 0.5">
-        <joint axis="1 0 0" type="hinge" range="-20 20"/>
-        <geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
-        <body pos="0 0.5 0">
-          <joint axis="1 0 0" type="hinge" range="-30 30"/>
-          <geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-</mujoco>

mjx/mujoco/mjx/test_data/cherry_pendulum.xml

@@ -1,20 +0,0 @@
-<mujoco model="cherry_pendulum">
-  <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
-  <default>
-    <geom contype="0" conaffinity="0"/>
-  </default>
-  <worldbody>
-    <body>
-      <joint axis="1 0 0" type="hinge" range="-45 45"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-      <body pos="0 0.5 0">
-        <joint axis="1 0 0" type="hinge"/>
-        <geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
-      </body>
-      <body pos="0 0.5 0">
-        <joint axis="1 0 0" type="hinge"/>
-        <geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
-      </body>
-    </body>
-  </worldbody>
-</mujoco>

mjx/mujoco/mjx/test_data/mixed_joint_pendulum.xml

@@ -1,23 +0,0 @@
-<mujoco model="revolute">
-  <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
-  <compiler autolimits="true"/>
-  <default>
-    <geom contype="0" conaffinity="0"/>
-    <joint damping="20"/>
-  </default>
-  <worldbody>
-    <body>
-      <joint axis="1 0 0" type="ball" range="0 10"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-      <body pos="0.3 0.4 0.5">
-        <joint axis="1 0 0" type="hinge" range="-20 20"/>
-        <joint axis="0 1 0" type="hinge" range="-20 20"/>
-        <geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
-        <body pos="0 0.5 0">
-          <joint axis="1 0 0" type="hinge" range="-30 30"/>
-          <geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-</mujoco>

mjx/mujoco/mjx/test_data/pendula.xml

@@ -0,0 +1,102 @@
+<!-- For validating dynamics of joints:
+
+* free, ball, slide, hinge joints
+* stacked joints (e.g. hinge + slide, ball + slide, etc)
+* n-link kinematic chains
+* limits, armature, damping
+-->
+<mujoco model="pendula">
+  <compiler autolimits="true"/>
+
+  <option timestep="0.02">
+    <flag contact="disable" />
+  </option>
+
+  <default>
+    <geom type="box" pos=".1 .2 .3" size=".1 .2 .3"/>
+    <joint damping="0.25" stiffness="0.1"/>
+  </default>
+
+  <worldbody>
+    <!-- a single free body -->
+    <body pos="0 0 0">
+      <freejoint/>
+      <geom/>
+    </body>
+
+    <!-- a single ball joint with a limit -->
+    <body pos="0.5 0 0">
+      <joint type="ball" range="0 35"/>
+      <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"/>
+      <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"/>
+      <geom/>
+    </body>
+
+    <!-- stacked joint: hinge + slide -->
+    <body pos="2.0 0 0">
+      <joint type="hinge" axis="0.1 0.2 0.3"/>
+      <joint 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 type="ball"/>
+      <geom/>
+    </body>
+
+    <!-- triple pendulum of hinges -->
+    <body pos="3.0 0 0">
+      <joint 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"/>
+        <geom/>
+        <body pos="0 -0.7 0">
+          <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <geom/>
+        </body>
+      </body>
+    </body>
+
+    <!-- cherry pendulum: two bodies attached to same parent body -->
+    <body pos="3.5 0 0">
+      <joint 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"/>
+        <geom/>
+      </body>
+      <body pos="0 -0.7 0">
+        <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+        <geom/>
+      </body>
+    </body>
+
+    <!-- falling pendulum -->
+    <body pos="4.0 0 0">
+      <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"/>
+        <geom/>
+        <body pos="0 -0.7 0">
+          <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <geom/>
+        </body>
+      </body>
+    </body>
+
+  </worldbody>
+</mujoco>

mjx/mujoco/mjx/test_data/single_pendulum.xml

@@ -1,13 +0,0 @@
-<mujoco model="pendulum">
-  <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
-  <worldbody>
-    <body>
-      <joint name="slider" axis="1 0 0" type="hinge"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-    </body>
-  </worldbody>
-  <!-- Tests that a single actuator doesn't get mangled in a physics step. -->
-  <actuator>
-    <motor name="slide" joint="slider" gear="10" ctrllimited="true" ctrlrange="-1 1"/>
-  </actuator>
-</mujoco>

mjx/mujoco/mjx/test_data/slide_pendulum.xml

@@ -1,20 +0,0 @@
-<mujoco model="slide_pendulum">
-  <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
-  <default>
-    <geom contype="0" conaffinity="0"/>
-  </default>
-  <worldbody>
-    <body>
-      <joint axis="1 0 0" type="hinge"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-      <body pos="0 0.5 0">
-        <joint axis="1 0 0" type="hinge"/>
-        <geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
-        <body pos="0 0.5 0">
-          <joint axis="1 0 0" type="slide"/>
-          <geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-</mujoco>

mjx/mujoco/mjx/test_data/triple_pendulum.xml

@@ -1,21 +0,0 @@
-<mujoco model="triple_pendulum">
-  <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
-  <compiler autolimits="true"/>
-  <default>
-    <geom contype="0" conaffinity="0"/>
-  </default>
-  <worldbody>
-    <body>
-      <joint axis="1 0 0" type="hinge" armature="0.01" range="-10 10"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-      <body pos="0 0.5 0">
-        <joint axis="1 0 0" type="hinge" armature="0.02" range="-20 20"/>
-        <geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
-        <body pos="0 0.5 0">
-          <joint axis="1 0 0" type="hinge" armature="0.03" range="-30 30"/>
-          <geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-</mujoco>

mjx/mujoco/mjx/test_data/triple_pendulum_free.xml

@@ -1,24 +0,0 @@
-<mujoco model="triple_pendulum_free">
-  <option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
-  <default>
-    <geom contype="0" conaffinity="0"/>
-  </default>
-  <worldbody>
-    <body>
-      <joint type="free"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-    </body>
-    <body>
-      <joint axis="1 0 0" type="hinge"/>
-      <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-      <body pos="0 0.5 0">
-        <joint axis="1 0 0" type="hinge"/>
-        <geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
-        <body pos="0 0.5 0">
-          <joint axis="1 0 0" type="hinge"/>
-          <geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-</mujoco>