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/__init__.py
+8
View File
@@ -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
+14 -22
View File
@@ -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')
d = mujoco.MjData(m)
d.qpos[2] = 0.0
mx = mjx.device_put(m)
ncon = collision_driver.ncon(mx)
self.assertEqual(ncon, 4)
m = test_util.load_test_file('constraints.xml')
ncon = collision_driver.ncon(m)
self.assertEqual(ncon, 16)
def test_disable_contact(self):
m = test_util.load_test_file('ant.xml')
d = mujoco.MjData(m)
d.qpos[2] = 0.0
m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT
mx = mjx.device_put(m)
ncon = collision_driver.ncon(mx)
m = test_util.load_test_file('constraints.xml')
m.opt.disableflags |= DisableBit.CONTACT
ncon = collision_driver.ncon(m)
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
+55 -122
View File
@@ -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):
err_msg = f'mismatch: {name} at step {step} in {fname}'
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
# tolerance for difference between MuJoCo and MJX constraint calculations,
# mostly due to float precision
_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)
# exclude convex.xml since convex contacts are not exactly equivalent
if fname == 'convex.xml':
return
m = test_util.load_test_file(fname)
m = test_util.load_test_file('constraints.xml')
d = mujoco.MjData(m)
mx = mjx.device_put(m)
dx = mjx.make_data(mx)
mujoco.mj_step(m, d, 100) # at 100 steps mix of active/inactive constraints
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
forward_jit_fn = jax.jit(mjx.forward)
# 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):
dx = dx.replace(qpos=jax.device_put(d.qpos), qvel=jax.device_put(d.qvel))
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)
dx = mjx.make_constraint(mx, dx)
nnz = dx.efc_J.any(axis=1)
_assert_eq(d.efc_J, dx.efc_J[nnz].reshape(-1), 'efc_J')
_assert_eq(d.efc_D, dx.efc_D[nnz], 'efc_D')
_assert_eq(d.efc_aref, dx.efc_aref[nnz], 'efc_aref')
_assert_eq(d.efc_frictionloss, dx.efc_frictionloss[nnz], 'efc_frictionloss')
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
mx = mjx.device_put(m)
k, *_ = constraint._kbi(mx, solimp, solref, pos)
self.assertEqual(k, 1 / (0.99**2 * (2 * m.opt.timestep) ** 2))
def test_disableconstraint(self):
m = test_util.load_test_file('ant.xml')
d = mujoco.MjData(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)
def test_disable_constraint(self):
m = test_util.load_test_file('constraints.xml')
m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONSTRAINT
ne, nf, nl, nc = mjx.count_constraints(m)
self.assertEqual(ne, 0)
self.assertEqual(nf, 0)
self.assertEqual(nl, 0)
self.assertEqual(nc, 0)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 0)
def test_disable_equality(self):
m = test_util.load_test_file('equality.xml')
d = mujoco.MjData(m)
m.opt.disableflags = m.opt.disableflags | DisableBit.EQUALITY
mx, dx = mjx.device_put(m), mjx.device_put(d)
dx = constraint.make_constraint(mx, dx)
self.assertEqual(dx.efc_J.shape[0], 0)
m = test_util.load_test_file('constraints.xml')
m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY
ne, nf, nl, nc = mjx.count_constraints(m)
self.assertEqual(ne, 0)
self.assertEqual(nf, 0)
self.assertEqual(nl, 2)
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')
d = mujoco.MjData(m)
d.qpos[2] = 0.0
mujoco.mj_forward(m, d)
m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONTACT
mx, dx = mjx.device_put(m), mjx.device_put(d)
dx = dx.tree_replace(
{'contact.frame': dx.contact.frame.reshape((-1, 3, 3))}
)
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)
m = test_util.load_test_file('constraints.xml')
m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
ne, nf, nl, nc = mjx.count_constraints(m)
self.assertEqual(ne, 10)
self.assertEqual(nf, 0)
self.assertEqual(nl, 2)
self.assertEqual(nc, 0)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 12) # only joint range, limit
if __name__ == '__main__':
mjx/mujoco/mjx/_src/device.py
+9
View File
@@ -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
+5 -5
View File
@@ -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
+12 -12
View File
@@ -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 = _velocity(m, d)
d = _actuation(m, d)
d = _acceleration(m, d)
d = fwd_position(m, d)
d = fwd_velocity(m, d)
d = fwd_actuation(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
+68 -77
View File
@@ -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):
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)
# tolerance for difference between MuJoCo and MJX forward calculations - mostly
# due to float precision
_TOLERANCE = 1e-5
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)
dx = mjx.make_data(mx)
forward_jit_fn = jax.jit(mjx.forward)
# apply some control and xfrc input
d.ctrl = np.array([-18, 0.59, 0.47])
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
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))
mx = mjx.put_model(m)
_assert_attr_eq(d, dx, 'qfrc_smooth', i, fname)
_assert_attr_eq(d, dx, 'qacc_smooth', i, fname)
# fwd_actuation
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(
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."""
np.random.seed(test_util.TEST_FILES.index(fname))
m = test_util.load_test_file(fname)
step_jit_fn = jax.jit(forward.step)
# fwd_accleration (fwd_position and fwd_velocity already tested elsewhere)
dx = jax.jit(mjx.fwd_acceleration)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qfrc_smooth')
_assert_attr_eq(d, dx, 'qacc_smooth')
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
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)
# apply some control and xfrc input
d.ctrl = np.array([-18, 0.59, 0.47])
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_step(m, d)
dx = step_jit_fn(mx, dx)
_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)
mx = mjx.put_model(m)
dx = jax.jit(mjx.step)(mx, mjx.put_data(m, d))
mujoco.mj_step(m, d)
_assert_attr_eq(d, dx, 'act')
_assert_attr_eq(d, dx, 'time')
_assert_attr_eq(d, dx, 'qvel')
_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
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)
d.qvel = np.array([0.2, -0.1])
mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero
mujoco.mj_forward(m, d)
mujoco.mj_step(m, d)
dx = step_jit_fn(mx, dx)
mx = mjx.put_model(m)
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, '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')
_assert_attr_eq(d, dx, 'qvel')
_assert_attr_eq(d, dx, 'qpos')
_assert_attr_eq(d, dx, 'act')
_assert_attr_eq(d, dx, 'time')
def test_disable_eulerdamp(self):
m = test_util.load_test_file('ant.xml')
m.opt.disableflags = m.opt.disableflags | DisableBit.EULERDAMP
m = test_util.load_test_file('pendula.xml')
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)
self.assertTrue((mx.dof_damping > 0).any())
dx = mjx.device_put(d)
dx = jax.jit(forward.forward)(mx, dx)
d.qvel[:] = 1.0
d.qacc[:] = 1.0
mx = mjx.put_model(m)
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
+2 -1
View File
@@ -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
+1 -1
View File
@@ -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
+43 -78
View File
@@ -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
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)
# tolerance for difference between MuJoCo and MJX passive calculations - mostly
# due to float precision
_TOLERANCE = 1e-7
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
m.jnt_stiffness = np.random.uniform(size=m.njnt)
m.dof_damping = np.random.uniform(size=m.nv)
def _assert_attr_eq(a, b, attr):
_assert_eq(getattr(a, attr), getattr(b, attr), attr)
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 = mjx.make_data(mx)
dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
_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):
qpos, qvel = d.qpos.copy(), d.qvel.copy()
mujoco.mj_step(m, d)
dx = passive_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
_assert_attr_eq(d, dx, 'qfrc_passive', i, fname)
m.opt.viscosity = 0.02
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')
@parameterized.parameters(
itertools.product(range(3), ('pendula.xml',))
)
def test_fluid(self, seed, fname):
np.random.seed(seed)
path = epath.resource_path('mujoco.mjx') / 'test_data'
path /= fname
m = mujoco.MjModel.from_xml_string(path.read_text())
m.opt.wind = np.array([0.03, 0.04, 0.05])
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')
# set density/viscosity/wind
m.opt.density = np.random.uniform()
m.opt.viscosity = np.random.uniform()
m.opt.wind = np.random.uniform()
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))
# test disable passive
mx = mx.tree_replace({'opt.disableflags': mjx.DisableBit.PASSIVE})
dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
np.testing.assert_allclose(dx.qfrc_passive, 0)
if __name__ == '__main__':
mjx/mujoco/mjx/_src/scan_test.py
+1 -1
View File
@@ -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
+10 -12
View File
@@ -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)
length = jp.dot(axis, gear[:3])[None]
axis, angle = math.quat_to_axis_angle(qpos)
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',
'uvvv',
'ujjq',
'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
+83 -148
View File
@@ -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
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)
# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
# due to float precision
_TOLERANCE = 5e-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)
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)
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):
qpos, qvel = d.qpos.copy(), d.qvel.copy()
mujoco.mj_step(m, d)
mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
# kinematics
dx = kinematics_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
_assert_attr_eq(d, dx, 'xanchor', i, fname)
_assert_attr_eq(d, dx, 'xaxis', i, fname)
_assert_attr_eq(d, dx, 'xpos', i, fname)
_assert_attr_eq(d, dx, 'xquat', i, fname)
_assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat', i, fname)
_assert_attr_eq(d, dx, 'xipos', i, fname)
_assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat', i, fname)
_assert_attr_eq(d, dx, 'geom_xpos', i, fname)
_assert_eq(
d.geom_xmat.reshape((-1, 3, 3)),
dx.geom_xmat,
'geom_xmat',
i,
fname,
)
# kinematics
dx = jax.jit(mjx.kinematics)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'xanchor')
_assert_attr_eq(d, dx, 'xaxis')
_assert_attr_eq(d, dx, 'xpos')
_assert_attr_eq(d, dx, 'xquat')
_assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat')
_assert_attr_eq(d, dx, 'xipos')
_assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat')
_assert_attr_eq(d, dx, 'geom_xpos')
_assert_eq(d.geom_xmat.reshape((-1, 3, 3)), dx.geom_xmat, 'geom_xmat')
_assert_attr_eq(d, dx, 'site_xpos')
_assert_eq(d.site_xmat.reshape((-1, 3, 3)), dx.site_xmat, 'site_xmat')
# com_pos
dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
_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
dx = com_pos_jit_fn(mx, dx)
_assert_attr_eq(d, dx, 'subtree_com', i, fname)
_assert_attr_eq(d, dx, 'cinert', i, fname)
_assert_attr_eq(d, dx, 'cdof', i, fname)
def test_mul_m(self):
m = test_util.load_test_file('pendula.xml')
d = mujoco.MjData(m)
# give the system a little kick to ensure we have non-identity rotations
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
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):
def test_disable_gravity(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)
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))
)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
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
+41 -96
View File
@@ -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):
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)
# tolerance for difference between MuJoCo and MJX constraint calculations,
# mostly due to float precision
_TOLERANCE = 5e-5
class Solver64Test(parameterized.TestCase):
"""Tests solvers at 64 bit precision."""
def _assert_eq(a, b, name, tol=_TOLERANCE):
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):
super().tearDown()
jax.config.update('jax_enable_x64', False)
def _assert_attr_eq(a, b, attr):
_assert_eq(getattr(a, attr), getattr(b, attr), attr)
@parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml')))
def test_cg(self, seed, fname):
"""Test mjx cg solver matches mujoco cg solver at 64 bit precision."""
f = epath.resource_path('mujoco.mjx') / 'test_data' / fname
m = mujoco.MjModel.from_xml_string(f.read_text())
class SolverTest(absltest.TestCase):
def test_solver(self):
"""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)
forward_jit_fn = jax.jit(mjx.forward)
dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
_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
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)
# 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)
# also test normal CG
m.opt.solver = mujoco.mjtSolver.mjSOL_CG
mujoco.mj_forward(m, d)
dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qacc_warmstart')
_assert_attr_eq(d, dx, 'qacc')
_assert_attr_eq(d, dx, 'qfrc_constraint')
_assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
# 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
+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)
mjx/mujoco/mjx/_src/test_util.py
+1 -3
View File
@@ -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
+2 -2
View File
@@ -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
+3 -4
View File
@@ -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)
dx = mjx.device_put(d)
mx = mjx.put_model(m)
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
+2 -2
View File
@@ -60,8 +60,8 @@ class TransmissionIntegrationTest(parameterized.TestCase):
d.qvel = np.random.random(m.nv)
# put on device
mx = mjx.device_put(m)
dx = mjx.device_put(d)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
mujoco.mj_transmission(m, d)
dx = transmission_jit_fn(mx, dx)
mjx/mujoco/mjx/test_data/ant.xml
-82
View File
@@ -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
+53
View File
@@ -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
-71
View File
@@ -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
-109
View File
@@ -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
+28 -15
View File
@@ -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 type="slide" axis="0.4 0.5 0.6" range="0 1"/>
<joint name="joint4" type="hinge" axis="0.1 0.2 0.3"/>
<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>