diff --git a/doc/changelog.rst b/doc/changelog.rst index fa35b8b8..7dddf213 100644 --- a/doc/changelog.rst +++ b/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 ^^^^^^^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 166b1892..313b43ff 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/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 diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 98448a7d..8ce1bfa6 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index d5bdf07b..40da667c 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/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)) - def test_forward(self, seed, fname): + @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.""" - if fname in ('equality.xml',): - return - - np.random.seed(seed) + np.random.seed(test_util.TEST_FILES.index(fname)) 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))) - def test_step(self, fname, integrator_type): + @parameterized.parameters( + 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 ( - '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) + np.random.seed(test_util.TEST_FILES.index(fname)) 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, 'qpos', i, test_name, atol=1e-2) - _assert_attr_eq(d, dx, 'act', i, test_name) - _assert_attr_eq(d, dx, 'time', i, test_name) + _assert_attr_eq(d, dx, 'qvel', i, fname, atol=1e-2) + _assert_attr_eq(d, dx, 'qpos', i, fname, atol=1e-2) + _assert_attr_eq(d, dx, 'act', i, fname) + _assert_attr_eq(d, dx, 'time', i, fname) + + def test_rk4(self): + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + + + + """) + 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') diff --git a/mjx/mujoco/mjx/_src/passive_test.py b/mjx/mujoco/mjx/_src/passive_test.py index 49264014..4a5026c8 100644 --- a/mjx/mujoco/mjx/_src/passive_test.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py index 2975a32f..de904e14 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/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 - key_typ_ids, key_body_ids = {}, {} - for body_id in np.arange(m.nbody, dtype=np.int32): - depth = depth_fn(body_id) + # group together bodies that will be processed together. grouping key: + # 1) the tree depth: parent bodies are processed first, so that they are + # available as carry input to child bodies (or reverse if reverse=True) + # 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 - if any(t in 'jqv' for t in in_types + out_types): - jnts = np.nonzero(typ_body_id['j'] == body_id)[0] - jnts_p = np.nonzero(typ_body_id['j'] == m.body_parentid[body_id])[0] - key = depth, tuple(m.jnt_type[jnts]) - parent_key = depth - 1, tuple(m.jnt_type[jnts_p]) - else: - key, parent_key = (depth, ()), (depth - 1, ()) + # map key => body id + key_body_ids = {} + for body_id in range(m.nbody): + parent_id = -1 + if body_id > 0: + parent_id = m.body_parentid[body_id] + depths[body_id] = 1 + depths[parent_id] + + # 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 - for t in set(in_types + out_types): - out = key_typ_ids.setdefault(key, {}) - id_ = body_id if t == 'b' else np.nonzero(typ_body_id[t] == body_id)[0] - id_ = np.expand_dims(id_, axis=0) - out[t] = np.concatenate((out[t], id_)) if t in out else id_ + # find parent keys of each key. a key may have multiple parents if the + # carry output keys of distinct parents are the same. e.g.: + # - depth 0 body 1 (slide joint) + # -- depth 1 body 1 (hinge joint) + # - depth 0 body 2 (ball joint) + # -- 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: - parent_key = key_parents[key] - y = key_y.get(parent_key) + elif key in key_parents: + ys = [key_y[p] for p in key_parents[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: - 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)] + f_args = [_take(arg, ids) for arg, ids in zip(args, key_in_take[key])] 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 - ys = [key_y[key] for key, _ in key_typ_ids] - f_ret_is_seq = isinstance(ys[0], (list, tuple)) - ys = ys if f_ret_is_seq else [[y] for y in ys] - ys = [ - [v if typ == 'b' else jp.concatenate(v) for v, typ in zip(y, out_types)] - for y in ys - ] - ys = jax.tree_map(lambda *x: jp.concatenate(x), *ys) + # concatenate ys, drop grouping dimensions, put back in order + y = [] + for i, typ in enumerate(out_types): + y_typ = [key_y[key] for key in keys] + if len(out_types) > 1: + y_typ = [y_[i] for y_ in y_typ] + if typ != 'b': + y_typ = jax.tree_map(jp.concatenate, y_typ) + 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 - 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] + y = y[0] if len(out_types) == 1 else y return y diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 71ff1d75..9e518672 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py index 9f677ff0..8d765644 100644 --- a/mjx/mujoco/mjx/_src/test_util.py +++ b/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', - 'single_pendulum.xml', - 'slide_pendulum.xml', - 'triple_pendulum.xml', - 'triple_pendulum_free.xml', + 'pendula.xml', ] _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity'] diff --git a/mjx/mujoco/mjx/test_data/ball_pendulum.xml b/mjx/mujoco/mjx/test_data/ball_pendulum.xml deleted file mode 100644 index 097b7f7b..00000000 --- a/mjx/mujoco/mjx/test_data/ball_pendulum.xml +++ /dev/null @@ -1,23 +0,0 @@ - - diff --git a/mjx/mujoco/mjx/test_data/cherry_pendulum.xml b/mjx/mujoco/mjx/test_data/cherry_pendulum.xml deleted file mode 100644 index ea9082de..00000000 --- a/mjx/mujoco/mjx/test_data/cherry_pendulum.xml +++ /dev/null @@ -1,20 +0,0 @@ - - diff --git a/mjx/mujoco/mjx/test_data/mixed_joint_pendulum.xml b/mjx/mujoco/mjx/test_data/mixed_joint_pendulum.xml deleted file mode 100644 index 16212c4d..00000000 --- a/mjx/mujoco/mjx/test_data/mixed_joint_pendulum.xml +++ /dev/null @@ -1,23 +0,0 @@ - - diff --git a/mjx/mujoco/mjx/test_data/pendula.xml b/mjx/mujoco/mjx/test_data/pendula.xml new file mode 100644 index 00000000..2363a3f8 --- /dev/null +++ b/mjx/mujoco/mjx/test_data/pendula.xml @@ -0,0 +1,102 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/mjx/mujoco/mjx/test_data/single_pendulum.xml b/mjx/mujoco/mjx/test_data/single_pendulum.xml deleted file mode 100644 index 98f5a0c6..00000000 --- a/mjx/mujoco/mjx/test_data/single_pendulum.xml +++ /dev/null @@ -1,13 +0,0 @@ - - diff --git a/mjx/mujoco/mjx/test_data/slide_pendulum.xml b/mjx/mujoco/mjx/test_data/slide_pendulum.xml deleted file mode 100644 index cb2312fb..00000000 --- a/mjx/mujoco/mjx/test_data/slide_pendulum.xml +++ /dev/null @@ -1,20 +0,0 @@ - - diff --git a/mjx/mujoco/mjx/test_data/triple_pendulum.xml b/mjx/mujoco/mjx/test_data/triple_pendulum.xml deleted file mode 100644 index e6cb7b14..00000000 --- a/mjx/mujoco/mjx/test_data/triple_pendulum.xml +++ /dev/null @@ -1,21 +0,0 @@ - - diff --git a/mjx/mujoco/mjx/test_data/triple_pendulum_free.xml b/mjx/mujoco/mjx/test_data/triple_pendulum_free.xml deleted file mode 100644 index 834eefac..00000000 --- a/mjx/mujoco/mjx/test_data/triple_pendulum_free.xml +++ /dev/null @@ -1,24 +0,0 @@ - -