Add MJX and bump version number to 3.0.0.

Co-authored-by: Baruch Tabanpour <btaba@google.com> PiperOrigin-RevId: 574327508 Change-Id: Ia9b62fbc929c6869dfcec87636b2e10d405a1060
2023-10-18 03:09:54 +01:00
parent 3f3d5a3b49
commit 8f9c690c85
117 changed files with 114430 additions and 81 deletions
@@ -0,0 +1,81 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for support."""
+
+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 support
+from mujoco.mjx._src import test_util
+import numpy as np
+
+
+class SupportTest(parameterized.TestCase):
+
+  @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'})
+  def test_jac(self, fname):
+    np.random.seed(0)
+
+    m = test_util.load_test_file(fname)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d)
+    mx = mjx.device_put(m)
+    dx = mjx.device_put(d)
+    point = np.random.randn(3)
+    body = np.random.choice(m.nbody)
+    jacp, jacr = jax.jit(support.jac)(mx, dx, point, body)
+
+    jacp_expected, jacr_expected = np.zeros((3, m.nv)), np.zeros((3, m.nv))
+    mujoco.mj_jac(m, d, jacp_expected, jacr_expected, point, body)
+    np.testing.assert_almost_equal(jacp, jacp_expected.T, 6)
+    np.testing.assert_almost_equal(jacr, jacr_expected.T, 6)
+
+  def test_xfrc_accumulate(self):
+    """Tests that xfrc_accumulate ouput matches mj_xfrcAccumulate."""
+    np.random.seed(0)
+
+    m = test_util.load_test_file('ant.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d)
+    mx = mjx.device_put(m)
+    dx = mjx.device_put(d)
+    self.assertFalse((dx.xipos == 0.0).all())
+
+    xfrc = np.random.rand(*dx.xfrc_applied.shape)
+
+    d.xfrc_applied[:] = xfrc
+    dx = dx.replace(xfrc_applied=jp.array(xfrc))
+
+    qfrc = jax.jit(support.xfrc_accumulate)(mx, dx)
+    qfrc_expected = np.zeros(m.nv)
+    for i in range(1, m.nbody):
+      mujoco.mj_applyFT(
+          m,
+          d,
+          d.xfrc_applied[i, :3],
+          d.xfrc_applied[i, 3:],
+          d.xipos[i],
+          i,
+          qfrc_expected,
+      )
+
+    np.testing.assert_almost_equal(qfrc, qfrc_expected, 6)
+
+
+if __name__ == '__main__':
+  absltest.main()