chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Cleanup MJX tests and migrate them to put_model/put_data.

Browse Source 
PiperOrigin-RevId: 588304729
Change-Id: I58075383e6eed64ae00cea305a568eba6465b0b0
This commit is contained in:
Erik Frey
2023-12-05 23:07:00 -08:00
committed by Copybara-Service
parent 899ba4b7ec
commit b3ccf67ebf
24 changed files with 446 additions and 868 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support_test.py
+5 -5
View File
@@ -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)
dx = mjx.device_put(d)
mx = mjx.put_model(m)
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)
dx = mjx.device_put(d)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
self.assertFalse((dx.xipos == 0.0).all())
xfrc = np.random.rand(*dx.xfrc_applied.shape)