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

Add pyink and isort config. Reformat.

Browse Source 
PiperOrigin-RevId: 704533915
Change-Id: I37e9fd51261bd166b725c7460fc65d02fed2b391
This commit is contained in:
Baruch Tabanpour
2024-12-09 21:10:00 -08:00
committed by Copybara-Service
parent 6f6244b739
commit f3b3024291
41 changed files with 895 additions and 480 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/bindings_test.py
+268 -142
View File
@@ -120,13 +120,17 @@ class MuJoCoBindingsTest(parameterized.TestCase):
xml_2 = rb"""<mujoco><geom name="box" type="box" size="1 1 1"/></mujoco>"""
xml_3 = rb"""<mujoco><geom name="ball" type="sphere" size="1"/></mujoco>"""
model = mujoco.MjModel.from_xml_string(
xml_1, {'model_.xml': xml_2, 'model__.xml': xml_3})
xml_1, {'model_.xml': xml_2, 'model__.xml': xml_3}
)
self.assertEqual(
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'plane'), 0)
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'plane'), 0
)
self.assertEqual(
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'box'), 1)
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'box'), 1
)
self.assertEqual(
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'ball'), 2)
mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'ball'), 2
)
def test_load_xml_repeated_asset_name(self):
# Assets aren't allowed to have the same filename (even if they have
@@ -139,23 +143,25 @@ class MuJoCoBindingsTest(parameterized.TestCase):
def test_can_read_array(self):
np.testing.assert_array_equal(
self.model.body_pos,
[[0, 0, 0], [0, 0, 0.1], [0, 0, 0], [0, 0, 0], [42.0, 0, 42.0]])
[[0, 0, 0], [0, 0, 0.1], [0, 0, 0], [0, 0, 0], [42.0, 0, 42.0]],
)
def test_can_set_array(self):
self.data.qpos = 0.12345
np.testing.assert_array_equal(
self.data.qpos, [0.12345]*len(self.data.qpos))
self.data.qpos, [0.12345] * len(self.data.qpos)
)
def test_array_is_a_view(self):
qpos_ref = self.data.qpos
self.data.qpos = 0.789
np.testing.assert_array_equal(
qpos_ref, [0.789]*len(self.data.qpos))
np.testing.assert_array_equal(qpos_ref, [0.789] * len(self.data.qpos))
# This test is disabled on PyPy as it uses sys.getrefcount
# However PyPy is not officially supported by MuJoCo
@absltest.skipIf(sys.implementation.name == 'pypy',
reason='requires sys.getrefcount')
@absltest.skipIf(
sys.implementation.name == 'pypy', reason='requires sys.getrefcount'
)
def test_array_keeps_struct_alive(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
qpos0 = model.qpos0
@@ -185,11 +191,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
def test_named_indexing_actuator_ctrl(self):
actuator_id = mujoco.mj_name2id(
self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, 'myactuator')
self.assertIs(self.data.actuator('myactuator'),
self.data.actuator(actuator_id))
self.assertIs(self.data.actuator('myactuator').ctrl,
self.data.actuator(actuator_id).ctrl)
self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, 'myactuator'
)
self.assertIs(
self.data.actuator('myactuator'), self.data.actuator(actuator_id)
)
self.assertIs(
self.data.actuator('myactuator').ctrl,
self.data.actuator(actuator_id).ctrl,
)
self.assertEqual(self.data.actuator('myactuator').ctrl.shape, (1,))
# Test that the indexer is returning a view into the underlying struct.
@@ -202,41 +212,49 @@ class MuJoCoBindingsTest(parameterized.TestCase):
def test_named_indexing_invalid_names_in_model(self):
with self.assertRaisesRegex(
KeyError,
r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]"):
r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]",
):
self.model.geom('badgeom')
def test_named_indexing_no_name_argument_in_model(self):
with self.assertRaisesRegex(
KeyError,
r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]"):
r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]",
):
self.model.joint()
def test_named_indexing_invalid_names_in_data(self):
with self.assertRaisesRegex(
KeyError,
r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]"):
r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]",
):
self.data.geom('badgeom')
def test_named_indexing_no_name_argument_in_data(self):
with self.assertRaisesRegex(
KeyError,
r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]"):
r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]",
):
self.data.jnt()
def test_named_indexing_invalid_index_in_model(self):
with self.assertRaisesRegex(
IndexError, r'Invalid index 3\. Valid indices from 0 to 2'):
IndexError, r'Invalid index 3\. Valid indices from 0 to 2'
):
self.model.geom(3)
with self.assertRaisesRegex(
IndexError, r'Invalid index -1\. Valid indices from 0 to 2'):
IndexError, r'Invalid index -1\. Valid indices from 0 to 2'
):
self.model.geom(-1)
def test_named_indexing_invalid_index_in_data(self):
with self.assertRaisesRegex(
IndexError, r'Invalid index 3\. Valid indices from 0 to 2'):
IndexError, r'Invalid index 3\. Valid indices from 0 to 2'
):
self.data.geom(3)
with self.assertRaisesRegex(
IndexError, r'Invalid index -1\. Valid indices from 0 to 2'):
IndexError, r'Invalid index -1\. Valid indices from 0 to 2'
):
self.data.geom(-1)
def test_named_indexing_geom_size(self):
@@ -267,45 +285,53 @@ class MuJoCoBindingsTest(parameterized.TestCase):
def test_named_indexing_ragged_qpos(self):
balljoint_id = mujoco.mj_name2id(
self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myball')
self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myball'
)
self.assertIs(self.data.joint('myball'), self.data.joint(balljoint_id))
self.assertIs(self.data.joint('myball').qpos,
self.data.joint(balljoint_id).qpos)
self.assertIs(
self.data.joint('myball').qpos, self.data.joint(balljoint_id).qpos
)
self.assertEqual(self.data.joint('myball').qpos.shape, (4,))
# Test that the indexer is returning a view into the underlying struct.
qpos_from_indexer = self.data.joint('myball').qpos
qpos_idx = self.model.jnt_qposadr[balljoint_id]
self.data.qpos[qpos_idx:qpos_idx+4] = [4, 5, 6, 7]
self.data.qpos[qpos_idx : qpos_idx + 4] = [4, 5, 6, 7]
np.testing.assert_array_equal(qpos_from_indexer, [4, 5, 6, 7])
self.data.joint('myball').qpos = [9, 8, 7, 6]
np.testing.assert_array_equal(self.data.qpos[qpos_idx:qpos_idx+4],
[9, 8, 7, 6])
np.testing.assert_array_equal(
self.data.qpos[qpos_idx : qpos_idx + 4], [9, 8, 7, 6]
)
def test_named_indexing_ragged2d_cdof(self):
freejoint_id = mujoco.mj_name2id(
self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myfree')
self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myfree'
)
self.assertIs(self.data.joint('myfree'), self.data.joint(freejoint_id))
self.assertIs(self.data.joint('myfree').cdof,
self.data.joint(freejoint_id).cdof)
self.assertIs(
self.data.joint('myfree').cdof, self.data.joint(freejoint_id).cdof
)
self.assertEqual(self.data.joint('myfree').cdof.shape, (6, 6))
# Test that the indexer is returning a view into the underlying struct.
cdof_from_indexer = self.data.joint('myfree').cdof
dof_idx = self.model.jnt_dofadr[freejoint_id]
self.data.cdof[dof_idx:dof_idx+6, :] = np.reshape(range(36), (6, 6))
np.testing.assert_array_equal(cdof_from_indexer,
np.reshape(range(36), (6, 6)))
self.data.cdof[dof_idx : dof_idx + 6, :] = np.reshape(range(36), (6, 6))
np.testing.assert_array_equal(
cdof_from_indexer, np.reshape(range(36), (6, 6))
)
self.data.joint('myfree').cdof = 42
np.testing.assert_array_equal(self.data.cdof[dof_idx:dof_idx+6], [[42]*6]*6)
np.testing.assert_array_equal(
self.data.cdof[dof_idx : dof_idx + 6], [[42] * 6] * 6
)
def test_named_indexing_repr_in_data(self):
expected_repr = '''<_MjDataGeomViews
expected_repr = """<_MjDataGeomViews
id: 1
name: 'mybox'
xmat: array([0., 0., 0., 0., 0., 0., 0., 0., 0.])
xpos: array([0., 0., 0.])
>'''
>"""
self.assertEqual(expected_repr, repr(self.data.geom('mybox')))
def test_named_indexing_body_repr_in_data(self):
@@ -328,8 +354,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
self.assertGreater(self.data._address, 0)
self.assertGreater(model2._address, 0)
self.assertGreater(data2._address, 0)
self.assertLen({self.model._address, self.data._address,
model2._address, data2._address}, 4)
self.assertLen(
{
self.model._address,
self.data._address,
model2._address,
data2._address,
},
4,
)
def test_mjmodel_can_read_and_write_opt(self):
self.assertEqual(self.model.opt.timestep, 0.002)
@@ -361,7 +394,9 @@ class MuJoCoBindingsTest(parameterized.TestCase):
def test_mjmodel_can_access_names_directly(self):
# mjModel offers direct access to names array, to allow usecases other than
# id2name
model_name = str(self.model.names[0:self.model.names.find(b'\0')], 'utf-8')
model_name = str(
self.model.names[0 : self.model.names.find(b'\0')], 'utf-8'
)
self.assertEqual(model_name, 'test')
start_index = self.model.name_geomadr[0]
@@ -402,15 +437,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
model_copy = copy.copy(self.model)
self.assertEqual(
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_JOINT, 0),
'myfree')
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_JOINT, 0), 'myfree'
)
self.assertEqual(
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0),
'myplane')
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0), 'myplane'
)
self.assertEqual(
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1),
'mybox')
mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1), 'mybox'
)
# Make sure it's a copy.
self.model.geom_size[1] = 0.5
@@ -420,7 +455,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
def test_mjdata_can_copy(self):
self.data.qpos = [0, 0, 0.1*np.sqrt(2) - 0.001,
np.cos(np.pi/8), np.sin(np.pi/8), 0, 0, 0,
1, 0, 0, 0]
1, 0, 0, 0] # fmt: skip
mujoco.mj_forward(self.model, self.data)
data_copy = copy.copy(self.data)
@@ -455,7 +490,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
contact_copy.append(copy.copy(self.data.contact[i]))
# Sort contacts in anticlockwise order
contact_copy = sorted(
contact_copy, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
contact_copy, key=lambda x: np.arctan2(x.pos[1], x.pos[0])
)
np.testing.assert_allclose(contact_copy[0].pos[:2], [-0.1, -0.1])
np.testing.assert_allclose(contact_copy[1].pos[:2], [0.1, -0.1])
np.testing.assert_allclose(contact_copy[2].pos[:2], [0.1, 0.1])
@@ -502,7 +538,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
# Sort contacts in anticlockwise order
sorted_contact = sorted(
contact, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
contact, key=lambda x: np.arctan2(x.pos[1], x.pos[0])
)
np.testing.assert_allclose(sorted_contact[0].pos[:2], [-0.1, -0.1])
np.testing.assert_allclose(sorted_contact[1].pos[:2], [0.1, -0.1])
np.testing.assert_allclose(sorted_contact[2].pos[:2], [0.1, 0.1])
@@ -589,7 +626,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
self.assertEqual(data2.ncon, 4)
self.assertEqual(data2.contact, self.data.contact)
self.data.qpos[3:7] = [np.cos(np.pi/8), np.sin(np.pi/8), 0, 0]
self.data.qpos[3:7] = [np.cos(np.pi / 8), np.sin(np.pi / 8), 0, 0]
self.data.qpos[2] *= (np.sqrt(2) - 1) * 0.1 - 1e-6
mujoco.mj_forward(self.model, self.data)
self.assertEqual(self.data.ncon, 2)
@@ -674,7 +711,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
def test_mju_rotVecQuat(self): # pylint: disable=invalid-name
vec = [1, 0, 0]
quat = [np.cos(np.pi/8), 0, 0, np.sin(np.pi/8)]
quat = [np.cos(np.pi / 8), 0, 0, np.sin(np.pi / 8)]
expected = np.array([1, 1, 0]) / np.sqrt(2)
# Check that the output argument works, and that the binding returns None.
@@ -722,7 +759,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
size = mujoco.mj_stateSize(self.model, spec)
state_bad_size = np.empty(size + 1, np.float64)
expected_message = ('state size should equal mj_stateSize(m, spec)')
expected_message = 'state size should equal mj_stateSize(m, spec)'
with self.assertRaisesWithLiteralMatch(TypeError, expected_message):
mujoco.mj_getState(self.model, self.data, state_bad_size, spec)
@@ -781,8 +818,9 @@ class MuJoCoBindingsTest(parameterized.TestCase):
mat = np.empty((3, 10), np.float64)
mujoco.mj_angmomMat(self.model, self.data, mat, 0)
np.testing.assert_almost_equal(mat @ self.data.qvel,
self.data.subtree_angmom[0, :])
np.testing.assert_almost_equal(
mat @ self.data.qvel, self.data.subtree_angmom[0, :]
)
def test_mj_jacSite(self): # pylint: disable=invalid-name
mujoco.mj_forward(self.model, self.data)
@@ -792,20 +830,22 @@ class MuJoCoBindingsTest(parameterized.TestCase):
jacp = np.empty((3, 10), np.float64)
mujoco.mj_jacSite(self.model, self.data, jacp, None, site_id)
expected_jacp = np.array(
[[0, 0, 0, 0, 0, 0, -1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
expected_jacp = np.array([
[0, 0, 0, 0, 0, 0, -1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
])
np.testing.assert_array_equal(jacp, expected_jacp)
# Call mj_jacSite with only jacr.
jacr = np.empty((3, 10), np.float64)
mujoco.mj_jacSite(self.model, self.data, None, jacr, site_id)
expected_jacr = np.array(
[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
expected_jacr = np.array([
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
])
np.testing.assert_array_equal(jacr, expected_jacr)
# Call mj_jacSite with both jacp and jacr.
@@ -818,12 +858,14 @@ class MuJoCoBindingsTest(parameterized.TestCase):
# Check that the jacp argument must have the right size.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, np.empty((3, 6), jacp.dtype), None, site_id)
self.model, self.data, np.empty((3, 6), jacp.dtype), None, site_id
)
# Check that the jacr argument must have the right size.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, None, np.empty((4, 7), jacr.dtype), site_id)
self.model, self.data, None, np.empty((4, 7), jacr.dtype), site_id
)
# The following two checks need to be done with fully initialized arrays,
# since pybind11 prints out the array's contents when generating TypeErrors.
@@ -832,12 +874,14 @@ class MuJoCoBindingsTest(parameterized.TestCase):
# Check that the jacp argument must have the right dtype.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, np.zeros(jacp.shape, int), None, site_id)
self.model, self.data, np.zeros(jacp.shape, int), None, site_id
)
# Check that the jacr argument must have the right dtype.
with self.assertRaises(TypeError):
mujoco.mj_jacSite(
self.model, self.data, None, np.zeros(jacr.shape, int), site_id)
self.model, self.data, None, np.zeros(jacr.shape, int), site_id
)
def test_docstrings(self): # pylint: disable=invalid-name
self.assertEqual(
@@ -845,13 +889,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
"""mj_versionString() -> str
Return the current version of MuJoCo as a null-terminated string.
""")
""",
)
self.assertEqual(
mujoco.mj_Euler.__doc__,
"""mj_Euler(m: mujoco._structs.MjModel, d: mujoco._structs.MjData) -> None
Euler integrator, semi-implicit in velocity.
""")
""",
)
def test_float_constant(self):
self.assertEqual(mujoco.mjMAXVAL, 1e10)
@@ -866,17 +912,19 @@ Euler integrator, semi-implicit in velocity.
self.assertLen(mujoco.mjVISSTRING, mujoco.mjtVisFlag.mjNVISFLAG)
self.assertLen(mujoco.mjRNDSTRING, mujoco.mjtRndFlag.mjNRNDFLAG)
self.assertEqual(mujoco.mjDISABLESTRING[11], 'Refsafe')
self.assertEqual(mujoco.mjVISSTRING[mujoco.mjtVisFlag.mjVIS_INERTIA],
('Inertia', '0', 'I'))
self.assertEqual(
mujoco.mjVISSTRING[mujoco.mjtVisFlag.mjVIS_INERTIA],
('Inertia', '0', 'I'),
)
def test_enum_values(self):
self.assertEqual(mujoco.mjtJoint.mjJNT_FREE, 0)
self.assertEqual(mujoco.mjtJoint.mjJNT_BALL, 1)
self.assertEqual(mujoco.mjtJoint.mjJNT_SLIDE, 2)
self.assertEqual(mujoco.mjtJoint.mjJNT_HINGE, 3)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1<<0)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1<<1)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1<<2)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1 << 0)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1 << 1)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1 << 2)
self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 7)
self.assertEqual(mujoco.mjtGeom.mjGEOM_PLANE, 0)
self.assertEqual(mujoco.mjtGeom.mjGEOM_HFIELD, 1)
@@ -899,8 +947,9 @@ Euler integrator, semi-implicit in velocity.
x = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k']
self.assertEqual(x[mujoco.mjtFrame.mjFRAME_WORLD], 'h')
self.assertEqual(
x[mujoco.mjtFrame.mjFRAME_GEOM:mujoco.mjtFrame.mjFRAME_CAMERA],
['c', 'd'])
x[mujoco.mjtFrame.mjFRAME_GEOM : mujoco.mjtFrame.mjFRAME_CAMERA],
['c', 'd'],
)
def test_enum_ops(self):
# Note: when modifying this test, make sure the enum value is an odd number
@@ -909,10 +958,12 @@ Euler integrator, semi-implicit in velocity.
self.assertEqual(mujoco.mjtFrame.mjFRAME_WORLD, 7.0)
self.assertEqual(7, mujoco.mjtFrame.mjFRAME_WORLD)
self.assertEqual(7.0, mujoco.mjtFrame.mjFRAME_WORLD)
self.assertEqual(mujoco.mjtFrame.mjFRAME_WORLD,
mujoco.mjtFrame.mjFRAME_WORLD)
self.assertNotEqual(mujoco.mjtFrame.mjFRAME_WORLD,
mujoco.mjtFrame.mjFRAME_NONE)
self.assertEqual(
mujoco.mjtFrame.mjFRAME_WORLD, mujoco.mjtFrame.mjFRAME_WORLD
)
self.assertNotEqual(
mujoco.mjtFrame.mjFRAME_WORLD, mujoco.mjtFrame.mjFRAME_NONE
)
self.assertEqual(-mujoco.mjtFrame.mjFRAME_WORLD, -7)
self.assertIsInstance(-mujoco.mjtFrame.mjFRAME_WORLD, int)
@@ -989,22 +1040,28 @@ Euler integrator, semi-implicit in velocity.
self.assertEqual(
mujoco.mjtDisableBit.mjDSBL_GRAVITY | mujoco.mjtDisableBit.mjDSBL_LIMIT,
72)
72,
)
self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE | 33, 33)
self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE & 33, 32)
self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE ^ 33, 1)
self.assertEqual(33 | mujoco.mjtDisableBit.mjDSBL_PASSIVE, 33)
self.assertEqual(33 & mujoco.mjtDisableBit.mjDSBL_PASSIVE, 32)
self.assertEqual(33 ^ mujoco.mjtDisableBit.mjDSBL_PASSIVE, 1)
self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL << 1,
mujoco.mjtDisableBit.mjDSBL_WARMSTART)
self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL >> 3,
mujoco.mjtDisableBit.mjDSBL_CONTACT)
self.assertEqual(
mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL << 1,
mujoco.mjtDisableBit.mjDSBL_WARMSTART,
)
self.assertEqual(
mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL >> 3,
mujoco.mjtDisableBit.mjDSBL_CONTACT,
)
def test_can_raise_error(self):
self.data.pstack = self.data.narena
with self.assertRaisesRegex(mujoco.FatalError,
r'\Amj_stackAlloc: insufficient memory:'):
with self.assertRaisesRegex(
mujoco.FatalError, r'\Amj_stackAlloc: insufficient memory:'
):
mujoco.mj_forward(self.model, self.data)
def test_mjcb_time(self):
@@ -1042,7 +1099,8 @@ Euler integrator, semi-implicit in velocity.
with self.assertRaises(TestError) as e:
mujoco.mj_forward(self.model, self.data)
self.assertEqual(
e.exception.args, ('string', (1, 2, 3), {'a': 1, 'b': 2}))
e.exception.args, ('string', (1, 2, 3), {'a': 1, 'b': 2})
)
# Should not raise now that we've cleared the callback.
mujoco.mj_forward(self.model, self.data)
@@ -1050,12 +1108,14 @@ Euler integrator, semi-implicit in velocity.
def test_mjcb_time_wrong_return_type(self):
with temporary_callback(mujoco.set_mjcb_time, lambda: 'string'):
with self.assertRaisesWithLiteralMatch(
TypeError, 'mjcb_time callback did not return a number'):
TypeError, 'mjcb_time callback did not return a number'
):
mujoco.mj_forward(self.model, self.data)
def test_mjcb_time_not_callable(self):
with self.assertRaisesWithLiteralMatch(
TypeError, 'callback is not an Optional[Callable]'):
TypeError, 'callback is not an Optional[Callable]'
):
mujoco.set_mjcb_time(1)
def test_mjcb_sensor(self):
@@ -1088,8 +1148,9 @@ Euler integrator, semi-implicit in velocity.
# This test is disabled on PyPy as it uses sys.getrefcount
# However PyPy is not officially supported by MuJoCo
@absltest.skipIf(sys.implementation.name == 'pypy',
reason='requires sys.getrefcount')
@absltest.skipIf(
sys.implementation.name == 'pypy', reason='requires sys.getrefcount'
)
def test_mjcb_control_not_leak_memory(self):
model_instances = []
data_instances = []
@@ -1110,8 +1171,9 @@ Euler integrator, semi-implicit in velocity.
# This test is disabled on PyPy as it uses sys.getrefcount
# However PyPy is not officially supported by MuJoCo
@absltest.skipIf(sys.implementation.name == 'pypy',
reason='requires sys.getrefcount')
@absltest.skipIf(
sys.implementation.name == 'pypy', reason='requires sys.getrefcount'
)
def test_mjdata_holds_ref_to_model(self):
data = mujoco.MjData(mujoco.MjModel.from_xml_string('<mujoco/>'))
model = data.model
@@ -1150,9 +1212,15 @@ Euler integrator, semi-implicit in velocity.
# When the scene is updated, geoms are added to the scene
# (ngeom is incremented)
mujoco.mj_forward(self.model, self.data)
mujoco.mjv_updateScene(self.model, self.data, mujoco.MjvOption(),
None, mujoco.MjvCamera(),
mujoco.mjtCatBit.mjCAT_ALL, scene)
mujoco.mjv_updateScene(
self.model,
self.data,
mujoco.MjvOption(),
None,
mujoco.MjvCamera(),
mujoco.mjtCatBit.mjCAT_ALL,
scene,
)
self.assertGreater(scene.ngeom, 0)
def test_mjv_scene_without_model(self):
@@ -1164,10 +1232,19 @@ Euler integrator, semi-implicit in velocity.
# mj_ray has tricky argument types
geomid = np.zeros(1, np.int32)
mujoco.mj_forward(self.model, self.data)
mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0,
geomid)
mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1],
[0, 0, 0, 0, 0, 0], 0, 0, geomid)
mujoco.mj_ray(
self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0, geomid
)
mujoco.mj_ray(
self.model,
self.data,
[0, 0, 0],
[0, 0, 1],
[0, 0, 0, 0, 0, 0],
0,
0,
geomid,
)
# Check that named arguments work
mujoco.mj_ray(
m=self.model,
@@ -1177,7 +1254,8 @@ Euler integrator, semi-implicit in velocity.
geomgroup=None,
flg_static=0,
bodyexclude=0,
geomid=geomid)
geomid=geomid,
)
def test_mj_multi_ray(self):
nray = 3
@@ -1201,14 +1279,13 @@ Euler integrator, semi-implicit in velocity.
geomid=geomid,
dist=dist,
nray=nray,
cutoff=mujoco.mjMAXVAL)
cutoff=mujoco.mjMAXVAL,
)
for i in range(0, 3):
self.assertEqual(
dist[i],
mujoco.mj_ray(
self.model, self.data, pnt, vec[i], None, 1, -1, geom1
),
mujoco.mj_ray(self.model, self.data, pnt, vec[i], None, 1, -1, geom1),
)
self.assertEqual(geomid[i], geom1)
self.assertEqual(geomid[i], geom_ex[i])
@@ -1217,16 +1294,28 @@ Euler integrator, semi-implicit in velocity.
def test_inverse_fd_none(self):
eps = 1e-6
flg_centered = 0
mujoco.mjd_inverseFD(self.model, self.data, eps, flg_centered,
None, None, None, None, None, None, None)
mujoco.mjd_inverseFD(
self.model,
self.data,
eps,
flg_centered,
None,
None,
None,
None,
None,
None,
None,
)
def test_geom_distance(self):
mujoco.mj_forward(self.model, self.data)
fromto = np.empty(6, np.float64)
dist = mujoco.mj_geomDistance(self.model, self.data, 0, 2, 200, fromto)
self.assertEqual(dist, 41.9)
np.testing.assert_array_equal(fromto,
np.array((42., 0., 0., 42., 0., 41.9)))
np.testing.assert_array_equal(
fromto, np.array((42.0, 0.0, 0.0, 42.0, 0.0, 41.9))
)
def test_inverse_fd(self):
eps = 1e-6
@@ -1238,8 +1327,19 @@ Euler integrator, semi-implicit in velocity.
ds_dv = np.zeros((self.model.nv, self.model.nsensordata))
ds_da = np.zeros((self.model.nv, self.model.nsensordata))
dm_dq = np.zeros((self.model.nv, self.model.nM))
mujoco.mjd_inverseFD(self.model, self.data, eps, flg_centered,
df_dq, df_dv, df_da, ds_dq, ds_dv, ds_da, dm_dq)
mujoco.mjd_inverseFD(
self.model,
self.data,
eps,
flg_centered,
df_dq,
df_dv,
df_da,
ds_dq,
ds_dv,
ds_da,
dm_dq,
)
self.assertGreater(np.linalg.norm(df_dq), eps)
self.assertGreater(np.linalg.norm(df_dv), eps)
self.assertGreater(np.linalg.norm(df_da), eps)
@@ -1272,15 +1372,17 @@ Euler integrator, semi-implicit in velocity.
n_total = 4
n_band = 1
n_dense = 1
dense = np.array([[1.0, 0, 0, 0.1],
[0, 2.0, 0, 0.2],
[0, 0, 3.0, 0.3],
[0.1, 0.2, 0.3, 4.0]])
band = np.zeros(n_band*(n_total-n_dense) + n_dense*n_total)
dense = np.array([
[1.0, 0, 0, 0.1],
[0, 2.0, 0, 0.2],
[0, 0, 3.0, 0.3],
[0.1, 0.2, 0.3, 4.0],
])
band = np.zeros(n_band * (n_total - n_dense) + n_dense * n_total)
mujoco.mju_dense2Band(band, dense, n_total, n_band, n_dense)
for i in range(4):
index = mujoco.mju_bandDiag(i, n_total, n_band, n_dense)
self.assertEqual(band[index], i+1)
self.assertEqual(band[index], i + 1)
dense2 = np.zeros((n_total, n_total))
flg_sym = 1
mujoco.mju_band2Dense(dense2, band, n_total, n_band, n_dense, flg_sym)
@@ -1288,20 +1390,22 @@ Euler integrator, semi-implicit in velocity.
vec = np.array([[2.0], [2.0], [3.0], [4.0]])
res = np.zeros_like(vec)
n_vec = 1
mujoco.mju_bandMulMatVec(res, band, vec,
n_total, n_band, n_dense, n_vec, flg_sym)
mujoco.mju_bandMulMatVec(
res, band, vec, n_total, n_band, n_dense, n_vec, flg_sym
)
np.testing.assert_array_equal(res, dense @ vec)
diag_add = 0
diag_mul = 0
mujoco.mju_cholFactorBand(band, n_total, n_band, n_dense,
diag_add, diag_mul)
mujoco.mju_cholFactorBand(
band, n_total, n_band, n_dense, diag_add, diag_mul
)
mujoco.mju_cholSolveBand(res, band, vec, n_total, n_band, n_dense)
np.testing.assert_almost_equal(res, np.linalg.solve(dense, vec))
def test_mju_box_qp(self):
n = 5
res = np.zeros(n)
r = np.zeros((n, n+7))
r = np.zeros((n, n + 7))
index = np.zeros(n, np.int32)
h = np.eye(n)
g = np.ones((n,))
@@ -1324,7 +1428,7 @@ Euler integrator, semi-implicit in velocity.
mat = np.linspace(0, 1, 16).reshape(4, 4)
res = np.empty((4, 4), np.float64)
mujoco.mju_symmetrize(res, mat)
np.testing.assert_array_equal(res, 0.5*(mat + mat.T))
np.testing.assert_array_equal(res, 0.5 * (mat + mat.T))
def test_mju_clip(self):
self.assertEqual(mujoco.mju_clip(1.5, 1.0, 2.0), 1.5)
@@ -1332,14 +1436,14 @@ Euler integrator, semi-implicit in velocity.
self.assertEqual(mujoco.mju_clip(1.5, 0.0, 1.0), 1.0)
def test_mju_mul_vec_mat_vec(self):
vec1 = np.array([1., 2., 3.])
vec2 = np.array([3., 2., 1.])
mat = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.)
vec1 = np.array([1.0, 2.0, 3.0])
vec2 = np.array([3.0, 2.0, 1.0])
mat = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.0)
def test_mju_dense_to_sparse(self):
mat = np.array([[0., 1., 0.], [2., 0., 3.]])
expected_vals = np.array([1., 2., 3.])
mat = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
expected_vals = np.array([1.0, 2.0, 3.0])
expected_rownnz = np.array([1, 2])
expected_rowadr = np.array([0, 1])
expected_colind = np.array([1, 0, 2])
@@ -1355,8 +1459,8 @@ Euler integrator, semi-implicit in velocity.
np.testing.assert_array_equal(col_ind, expected_colind)
def test_mju_sparse_to_dense(self):
expected = np.array([[0., 1., 0.], [2., 0., 3.]])
mat = np.array((1., 2., 3.))
expected = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
mat = np.array((1.0, 2.0, 3.0))
rownnz = np.array([1, 2])
rowadr = np.array([0, 1])
colind = np.array([1, 0, 2])
@@ -1366,10 +1470,10 @@ Euler integrator, semi-implicit in velocity.
def test_mju_euler_to_quat(self):
quat = np.zeros(4)
euler = np.array([0, np.pi/2, 0])
euler = np.array([0, np.pi / 2, 0])
seq = 'xyz'
mujoco.mju_euler2Quat(quat, euler, seq)
expected_quat = np.array([np.sqrt(0.5), 0, np.sqrt(0.5), 0.])
expected_quat = np.array([np.sqrt(0.5), 0, np.sqrt(0.5), 0.0])
np.testing.assert_almost_equal(quat, expected_quat)
error = 'mju_euler2Quat: seq must contain exactly 3 characters'
@@ -1377,7 +1481,7 @@ Euler integrator, semi-implicit in velocity.
mujoco.mju_euler2Quat(quat, euler, 'xy')
with self.assertRaisesWithLiteralMatch(mujoco.FatalError, error):
mujoco.mju_euler2Quat(quat, euler, 'xyzy')
error = 'mju_euler2Quat: seq[2] is \'p\', should be one of x, y, z, X, Y, Z'
error = "mju_euler2Quat: seq[2] is 'p', should be one of x, y, z, X, Y, Z"
with self.assertRaisesWithLiteralMatch(mujoco.FatalError, error):
mujoco.mju_euler2Quat(quat, euler, 'xYp')
@@ -1396,8 +1500,16 @@ Euler integrator, semi-implicit in velocity.
mujoco.mj_step(self.model, self.data)
data2 = pickle.loads(pickle.dumps(self.data))
attr_to_compare = (
'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
'warning', 'energy', 'contact', 'efc_J'
'time',
'qpos',
'qvel',
'qacc',
'xpos',
'mocap_pos',
'warning',
'energy',
'contact',
'efc_J',
)
self._assert_attributes_equal(data2, self.data, attr_to_compare)
for _ in range(10):
@@ -1410,8 +1522,16 @@ Euler integrator, semi-implicit in velocity.
mujoco.mj_step(self.model, self.data)
data2 = pickle.loads(pickle.dumps(self.data))
attr_to_compare = (
'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
'warning', 'energy', 'contact', 'efc_J'
'time',
'qpos',
'qvel',
'qacc',
'xpos',
'mocap_pos',
'warning',
'energy',
'contact',
'efc_J',
)
self._assert_attributes_equal(data2, self.data, attr_to_compare)
for _ in range(10):
@@ -1422,7 +1542,10 @@ Euler integrator, semi-implicit in velocity.
def test_pickle_mjmodel(self):
model2 = pickle.loads(pickle.dumps(self.model))
attr_to_compare = (
'nq', 'nmat', 'body_pos', 'names',
'nq',
'nmat',
'body_pos',
'names',
)
self._assert_attributes_equal(model2, self.model, attr_to_compare)
@@ -1506,8 +1629,11 @@ Euler integrator, semi-implicit in velocity.
else:
self.assertEqual(actual_value, expected_value)
except AssertionError as e:
self.fail("Attribute '{}' differs from expected value: {}".format(
name, str(e)))
self.fail(
"Attribute '{}' differs from expected value: {}".format(
name, str(e)
)
)
if __name__ == '__main__':
python/mujoco/memory_leak_test.py
+2 -1
View File
@@ -57,6 +57,7 @@ class MemoryLeakTest(absltest.TestCase):
soft = -1
try:
import resource # pylint: disable=g-import-not-at-top
soft, hard = resource.getrlimit(resource.RLIMIT_AS)
resource.setrlimit(resource.RLIMIT_AS, (limit_in_bytes, hard))
except (ImportError, ValueError):
@@ -65,5 +66,5 @@ class MemoryLeakTest(absltest.TestCase):
return soft
if __name__ == '__main__':
if __name__ == "__main__":
absltest.main()
python/mujoco/minimize.py
+18 -12
View File
@@ -209,7 +209,7 @@ def least_squares(
# Decrease mu agressively: sequential decreases grow exponentially.
def decrease_mu(mu, n_reduc):
dmu = (1/mu_factor) ** (2**n_reduc)
dmu = (1 / mu_factor) ** (2**n_reduc)
mu = 0.0 if mu * dmu < mu_min else mu * dmu
n_reduc += 1
return mu, n_reduc
@@ -427,7 +427,6 @@ def jacobian_fd(
Returns:
jac: Jacobian of the residual at x.
n_res: updated number of residual evaluations (add x.size).
"""
n = x.size
if bounds is None:
@@ -438,7 +437,7 @@ def jacobian_fd(
xh = x + np.diag(eps_vec)
rh = residual(xh)
jac = (rh - r) / eps_vec
return jac, n_res+n
return jac, n_res + n
def check_jacobian(
@@ -467,14 +466,15 @@ def check_jacobian(
Returns:
n_res: updated number of residual evaluations.
"""
jac_fd, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds)
denom = np.abs(jac).sum() + np.abs(jac_fd).sum() + 1e-8
rel_diff = np.abs(jac - jac_fd) / denom
if np.any(rel_diff > 1e-5):
raise ValueError(f'User-provided {name} does not match finite-differences '
'to a relative tolerance of 1e-5.')
raise ValueError(
f'User-provided {name} does not match finite-differences '
'to a relative tolerance of 1e-5.'
)
print(f'User-provided {name} matches finite-differences.', file=output)
return n_res
@@ -489,8 +489,8 @@ def check_norm(
Args:
r: residual vector.
norm: Norm function returning either the norm scalar or its gradient
and Gauss-Newton Hessian.
norm: Norm function returning either the norm scalar or its gradient and
Gauss-Newton Hessian.
eps: finite-difference step size.
output: Optional file or StringIO to which to print messages.
"""
@@ -506,12 +506,16 @@ def check_norm(
# Check that Hessian is positive-definite.
if np.any(np.linalg.eigvals(n_h) < 0):
h_min = np.min(np.linalg.eigvals(n_h))
raise ValueError('User-provided norm Hessian is not positive definite. '
f'Minimum eigenvalue is {h_min:<.4g}')
raise ValueError(
'User-provided norm Hessian is not positive definite. '
f'Minimum eigenvalue is {h_min:<.4g}'
)
# Local function returning norm values (vectorized).
def norm_vec(v):
norms = [np.atleast_2d(norm.value(v[:, i:i+1])) for i in range(v.shape[1])]
norms = [
np.atleast_2d(norm.value(v[:, i : i + 1])) for i in range(v.shape[1])
]
return np.hstack(norms)
# Check the norm gradient.
@@ -519,7 +523,9 @@ def check_norm(
# Local function returning norm gradients (vectorized).
def grad_vec(v):
gradients = [norm.grad_hess(v[:, i:i+1], eye)[0] for i in range(v.shape[1])]
gradients = [
norm.grad_hess(v[:, i : i + 1], eye)[0] for i in range(v.shape[1])
]
return np.hstack(gradients)
# Check the norm Hessian.
python/mujoco/minimize_test.py
+70 -31
View File
@@ -56,8 +56,9 @@ class MinimizeTest(absltest.TestCase):
x0 = np.array((0.0, 0.0))
out = io.StringIO()
x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out,
check_derivatives=True)
x, _ = minimize.least_squares(
x0, residual, jacobian=jacobian, output=out, check_derivatives=True
)
expected_x = np.array((1.0, 1.0))
np.testing.assert_array_almost_equal(x, expected_x)
self.assertIn('norm(dx) < tol', out.getvalue())
@@ -67,9 +68,15 @@ class MinimizeTest(absltest.TestCase):
def bad_jacobian(x, r):
del r # Unused.
return np.array([[-1, 0], [-20 * x[0, 0], 15]])
with self.assertRaisesRegex(ValueError, r'\bJacobian does not match\b'):
minimize.least_squares(x0, residual, jacobian=bad_jacobian, output=out,
check_derivatives=True)
minimize.least_squares(
x0,
residual,
jacobian=bad_jacobian,
output=out,
check_derivatives=True,
)
def test_max_iter(self) -> None:
dim = 20 # High-D Rosenbrock
@@ -98,13 +105,16 @@ class MinimizeTest(absltest.TestCase):
x0 = np.array((0.0, 0.0))
expected_x = np.array((1.0, 1.0))
bounds_types = {'inbounds': [np.array((-2.0, -2.0)), np.array((2.0, 2.0))],
'onlower': [np.array((-2.0, 2.0)), np.array((0.5, 3.0))],
'onupper': [np.array((-2.0, -2.0)), np.array((0.5, 2.0))]}
bounds_types = {
'inbounds': [np.array((-2.0, -2.0)), np.array((2.0, 2.0))],
'onlower': [np.array((-2.0, 2.0)), np.array((0.5, 3.0))],
'onupper': [np.array((-2.0, -2.0)), np.array((0.5, 2.0))],
}
# In bounds finds true minimum.
x, _ = minimize.least_squares(x0, residual, bounds=bounds_types['inbounds'],
output=out)
x, _ = minimize.least_squares(
x0, residual, bounds=bounds_types['inbounds'], output=out
)
np.testing.assert_array_almost_equal(x, expected_x)
self.assertIn('norm(dx) < tol', out.getvalue())
@@ -157,8 +167,9 @@ class MinimizeTest(absltest.TestCase):
print(f'Hello iteration {len(trace)}!', file=out)
x0 = np.array((0.0, 0.0))
x, _ = minimize.least_squares(x0, residual, output=out,
iter_callback=iter_callback)
x, _ = minimize.least_squares(
x0, residual, output=out, iter_callback=iter_callback
)
expected_x = np.array((1.0, 1.0))
np.testing.assert_array_almost_equal(x, expected_x)
self.assertIn('Hello iteration 3!', out.getvalue())
@@ -170,11 +181,12 @@ class MinimizeTest(absltest.TestCase):
p = 0.01 # Smoothing radius for smooth-L2 norm.
class SmoothL2(minimize.Norm):
def value(self, r):
return np.sqrt((r.T @ r).item() + p*p) - p
return np.sqrt((r.T @ r).item() + p * p) - p
def grad_hess(self, r, proj):
s = np.sqrt((r.T @ r).item() + p*p)
s = np.sqrt((r.T @ r).item() + p * p)
y_r = r / s
grad = proj.T @ y_r
y_rr = (np.eye(r.size) - y_r @ y_r.T) / s
@@ -183,8 +195,9 @@ class MinimizeTest(absltest.TestCase):
out = io.StringIO()
x0 = np.array((0.0, 0.0))
x, _ = minimize.least_squares(x0, residual, norm=SmoothL2(), output=out,
check_derivatives=True)
x, _ = minimize.least_squares(
x0, residual, norm=SmoothL2(), output=out, check_derivatives=True
)
expected_x = np.array((1.0, 1.0))
np.testing.assert_array_almost_equal(x, expected_x)
self.assertIn('norm(dx) < tol', out.getvalue())
@@ -192,11 +205,12 @@ class MinimizeTest(absltest.TestCase):
self.assertIn('User-provided norm Hessian matches', out.getvalue())
class SmoothL2BadGrad(minimize.Norm):
def value(self, r):
return np.sqrt((r.T @ r).item() + p*p) - p
return np.sqrt((r.T @ r).item() + p * p) - p
def grad_hess(self, r, proj):
s = np.sqrt((r.T @ r).item() + p*p)
s = np.sqrt((r.T @ r).item() + p * p)
y_r = r / s
grad = proj.T @ (y_r + 0.001) # 0.001 is erronous.
y_rr = (np.eye(r.size) - y_r @ y_r.T) / s
@@ -204,15 +218,21 @@ class MinimizeTest(absltest.TestCase):
return grad, hess
with self.assertRaisesRegex(ValueError, r'\bgradient does not match\b'):
minimize.least_squares(x0, residual, norm=SmoothL2BadGrad(), output=out,
check_derivatives=True)
minimize.least_squares(
x0,
residual,
norm=SmoothL2BadGrad(),
output=out,
check_derivatives=True,
)
class SmoothL2BadHess(minimize.Norm):
def value(self, r):
return np.sqrt((r.T @ r).item() + p*p) - p
return np.sqrt((r.T @ r).item() + p * p) - p
def grad_hess(self, r, proj):
s = np.sqrt((r.T @ r).item() + p*p)
s = np.sqrt((r.T @ r).item() + p * p)
y_r = r / s
grad = proj.T @ y_r
y_rr = (1.001 * np.eye(r.size) - y_r @ y_r.T) / s # 1.001 is erronous.
@@ -220,15 +240,21 @@ class MinimizeTest(absltest.TestCase):
return grad, hess
with self.assertRaisesRegex(ValueError, r'\bHessian does not match\b'):
minimize.least_squares(x0, residual, norm=SmoothL2BadHess(), output=out,
check_derivatives=True)
minimize.least_squares(
x0,
residual,
norm=SmoothL2BadHess(),
output=out,
check_derivatives=True,
)
class SmoothL2AsymHess(minimize.Norm):
def value(self, r):
return np.sqrt((r.T @ r).item() + p*p) - p
return np.sqrt((r.T @ r).item() + p * p) - p
def grad_hess(self, r, proj):
s = np.sqrt((r.T @ r).item() + p*p)
s = np.sqrt((r.T @ r).item() + p * p)
y_r = r / s
grad = proj.T @ y_r
y_rr = (np.eye(r.size) - (y_r + 0.0001) @ y_r.T) / s
@@ -236,15 +262,21 @@ class MinimizeTest(absltest.TestCase):
return grad, hess
with self.assertRaisesRegex(ValueError, r'\bnot symmetric\b'):
minimize.least_squares(x0, residual, norm=SmoothL2AsymHess(), output=out,
check_derivatives=True)
minimize.least_squares(
x0,
residual,
norm=SmoothL2AsymHess(),
output=out,
check_derivatives=True,
)
class SmoothL2NegHess(minimize.Norm):
def value(self, r):
return np.sqrt((r.T @ r).item() + p*p) - p
return np.sqrt((r.T @ r).item() + p * p) - p
def grad_hess(self, r, proj):
s = np.sqrt((r.T @ r).item() + p*p)
s = np.sqrt((r.T @ r).item() + p * p)
y_r = r / s
grad = proj.T @ y_r
y_rr = -(np.eye(r.size) - y_r @ y_r.T) / s # Negative-definite.
@@ -252,7 +284,14 @@ class MinimizeTest(absltest.TestCase):
return grad, hess
with self.assertRaisesRegex(ValueError, r'\bnot positive definite\b'):
minimize.least_squares(x0, residual, norm=SmoothL2NegHess(), output=out,
check_derivatives=True)
minimize.least_squares(
x0,
residual,
norm=SmoothL2NegHess(),
output=out,
check_derivatives=True,
)
if __name__ == '__main__':
absltest.main()
python/mujoco/msh2obj_test.py
+3 -1
View File
@@ -63,7 +63,8 @@ class MshTest(absltest.TestCase):
obj = msh2obj.msh_to_obj(msh_path)
obj_model = mujoco.MjModel.from_xml_string(
_XML, {"abdomen_1_body.obj": obj.encode()})
_XML, {"abdomen_1_body.obj": obj.encode()}
)
for field in _MESH_FIELDS:
np.testing.assert_allclose(
@@ -73,5 +74,6 @@ class MshTest(absltest.TestCase):
err_msg=f"Field {field} does not match between msh and obj models.",
)
if __name__ == "__main__":
absltest.main()
python/mujoco/render_test.py
+22 -8
View File
@@ -19,8 +19,9 @@ import mujoco
import numpy as np
@absltest.skipUnless(hasattr(mujoco, 'GLContext'),
'MuJoCo rendering is disabled')
@absltest.skipUnless(
hasattr(mujoco, 'GLContext'), 'MuJoCo rendering is disabled'
)
class MuJoCoRenderTest(absltest.TestCase):
def setUp(self):
@@ -48,8 +49,14 @@ class MuJoCoRenderTest(absltest.TestCase):
scene = mujoco.MjvScene(self.model, maxgeom=0)
mujoco.mjv_updateScene(
self.model, self.data, mujoco.MjvOption(), mujoco.MjvPerturb(),
mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
self.model,
self.data,
mujoco.MjvOption(),
mujoco.MjvPerturb(),
mujoco.MjvCamera(),
mujoco.mjtCatBit.mjCAT_ALL,
scene,
)
context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, context)
@@ -62,7 +69,7 @@ class MuJoCoRenderTest(absltest.TestCase):
mujoco.mjr_rectangle(blue_rect, 0, 0, 1, 1)
expected_upside_down_image = np.zeros((480, 640, 3), dtype=np.uint8)
expected_upside_down_image[67:67+123, 56:56+234, 2] = 255
expected_upside_down_image[67 : 67 + 123, 56 : 56 + 234, 2] = 255
upside_down_image = np.empty((480, 640, 3), dtype=np.uint8)
mujoco.mjr_readPixels(upside_down_image, None, full_rect, context)
@@ -71,7 +78,8 @@ class MuJoCoRenderTest(absltest.TestCase):
# Check that mjr_readPixels can accept a flattened array.
upside_down_image[:] = 0
mujoco.mjr_readPixels(
np.reshape(upside_down_image, -1), None, full_rect, context)
np.reshape(upside_down_image, -1), None, full_rect, context
)
np.testing.assert_array_equal(upside_down_image, expected_upside_down_image)
context.free()
@@ -81,8 +89,14 @@ class MuJoCoRenderTest(absltest.TestCase):
scene = mujoco.MjvScene(self.model, maxgeom=0)
mujoco.mjv_updateScene(
self.model, self.data, mujoco.MjvOption(), None,
mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
self.model,
self.data,
mujoco.MjvOption(),
None,
mujoco.MjvCamera(),
mujoco.mjtCatBit.mjCAT_ALL,
scene,
)
context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, context)
python/mujoco/renderer.py
+13 -11
View File
@@ -32,7 +32,7 @@ class Renderer:
model: _structs.MjModel,
height: int = 240,
width: int = 320,
max_geom: int = 10000
max_geom: int = 10000,
) -> None:
"""Initializes a new `Renderer`.
@@ -43,6 +43,7 @@ class Renderer:
max_geom: Optional integer specifying the maximum number of geoms that can
be rendered in the same scene. If None this will be chosen automatically
based on the estimated maximum number of renderable geoms in the model.
Raises:
ValueError: If `camera_id` is outside the valid range, or if `width` or
`height` exceed the dimensions of MuJoCo's offscreen framebuffer.
@@ -220,9 +221,7 @@ the clause:
# Convert 3-channel uint8 to 1-channel uint32.
image3 = out.astype(np.uint32)
segimage = (
image3[:, :, 0]
+ image3[:, :, 1] * (2**8)
+ image3[:, :, 2] * (2**16)
image3[:, :, 0] + image3[:, :, 1] * (2**8) + image3[:, :, 2] * (2**16)
)
# Remap segid to 2-channel (object ID, object type) pair.
# Seg ID 0 is background -- will be remapped to (-1, -1).
@@ -251,15 +250,15 @@ the clause:
self,
data: _structs.MjData,
camera: Union[int, str, _structs.MjvCamera] = -1,
scene_option: Optional[_structs.MjvOption] = None
):
scene_option: Optional[_structs.MjvOption] = None,
):
"""Updates geometry used for rendering.
Args:
data: An instance of `MjData`.
camera: An instance of `MjvCamera`, a string or an integer
scene_option: A custom `MjvOption` instance to use to render
the scene instead of the default.
scene_option: A custom `MjvOption` instance to use to render the scene
instead of the default.
Raises:
ValueError: If `camera_id` is outside the valid range, or if camera does
@@ -274,8 +273,10 @@ the clause:
if camera_id == -1:
raise ValueError(f'The camera "{camera}" does not exist.')
if camera_id < -1 or camera_id >= self._model.ncam:
raise ValueError(f'The camera id {camera_id} is out of'
f' range [-1, {self._model.ncam}).')
raise ValueError(
f'The camera id {camera_id} is out of'
f' range [-1, {self._model.ncam}).'
)
# Render camera.
camera = _structs.MjvCamera()
@@ -295,7 +296,8 @@ the clause:
data,
scene_option,
None,
camera, _enums.mjtCatBit.mjCAT_ALL.value,
camera,
_enums.mjtCatBit.mjCAT_ALL.value,
self._scene,
)
python/mujoco/renderer_test.py
+4 -2
View File
@@ -20,9 +20,11 @@ import mujoco
import numpy as np
@absltest.skipUnless(hasattr(mujoco, 'GLContext'),
'MuJoCo rendering is disabled')
@absltest.skipUnless(
hasattr(mujoco, 'GLContext'), 'MuJoCo rendering is disabled'
)
class MuJoCoRendererTest(parameterized.TestCase):
def test_renderer_unknown_camera_name(self):
xml = """
<mujoco>
python/mujoco/rollout.py
+70 -43
View File
@@ -23,17 +23,19 @@ import numpy as np
from numpy import typing as npt
def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
data: mujoco.MjData,
initial_state: npt.ArrayLike,
control: Optional[npt.ArrayLike] = None,
*, # require subsequent arguments to be named
control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
skip_checks: bool = False,
nstep: Optional[int] = None,
initial_warmstart: Optional[npt.ArrayLike] = None,
state: Optional[npt.ArrayLike] = None,
sensordata: Optional[npt.ArrayLike] = None):
def rollout(
model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
data: mujoco.MjData,
initial_state: npt.ArrayLike,
control: Optional[npt.ArrayLike] = None,
*, # require subsequent arguments to be named
control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
skip_checks: bool = False,
nstep: Optional[int] = None,
initial_warmstart: Optional[npt.ArrayLike] = None,
state: Optional[npt.ArrayLike] = None,
sensordata: Optional[npt.ArrayLike] = None,
):
"""Rolls out open-loop trajectories from initial states, get subsequent states and sensor values.
Python wrapper for rollout.cc, see documentation therein.
@@ -66,15 +68,24 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
Raises:
ValueError: bad shapes or sizes.
"""
""" # fmt: skip
# skip_checks shortcut:
# don't infer nroll/nstep
# don't support singleton expansion
# don't allocate output arrays
# just call rollout and return
if skip_checks:
_rollout.rollout(model, data, nstep, control_spec, initial_state,
initial_warmstart, control, state, sensordata)
_rollout.rollout(
model,
data,
nstep,
control_spec,
initial_state,
initial_warmstart,
control,
state,
sensordata,
)
return state, sensordata
if not isinstance(model, mujoco.MjModel):
@@ -92,17 +103,16 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
initial_warmstart=initial_warmstart,
control=control,
state=state,
sensordata=sensordata)
sensordata=sensordata,
)
# check number of dimensions
_check_number_of_dimensions(2,
initial_state=initial_state,
initial_warmstart=initial_warmstart)
_check_number_of_dimensions(3,
control=control,
state=state,
sensordata=sensordata)
_check_number_of_dimensions(
2, initial_state=initial_state, initial_warmstart=initial_warmstart
)
_check_number_of_dimensions(
3, control=control, state=state, sensordata=sensordata
)
# ensure 2D, make contiguous, row-major (C ordering)
initial_state = _ensure_2d(initial_state)
@@ -114,38 +124,46 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
sensordata = _ensure_3d(sensordata)
# infer nroll, check for incompatibilities
nroll = _infer_dimension(0, 1,
initial_state=initial_state,
initial_warmstart=initial_warmstart,
control=control,
state=state,
sensordata=sensordata)
nroll = _infer_dimension(
0,
1,
initial_state=initial_state,
initial_warmstart=initial_warmstart,
control=control,
state=state,
sensordata=sensordata,
)
if isinstance(model, list) and nroll == 1:
nroll = len(model)
if isinstance(model, list) and len(model) != nroll:
raise ValueError(f'nroll inferred as {nroll} '
f'but model is length {len(model)}')
raise ValueError(
f'nroll inferred as {nroll} but model is length {len(model)}'
)
elif not isinstance(model, list):
model = [model] # Use a length 1 list to simplify code below
model = [model] # Use a length 1 list to simplify code below
# infer nstep, check for incompatibilities
nstep = _infer_dimension(1, nstep or 1,
control=control,
state=state,
sensordata=sensordata)
nstep = _infer_dimension(
1, nstep or 1, control=control, state=state, sensordata=sensordata
)
# get nstate/ncontrol/nv/nsensordata
# check that they are equal across models
nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
nstate = mujoco.mj_stateSize(
model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value
)
ncontrol = mujoco.mj_stateSize(model[0], control_spec)
nv = model[0].nv
nsensordata = model[0].nsensordata
for m in model[1:]:
if (nstate != mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
if (
nstate
!= mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
or ncontrol != mujoco.mj_stateSize(m, control_spec)
or nv != m.nv
or nsensordata != m.nsensordata):
or nsensordata != m.nsensordata
):
raise ValueError('models are not compatible')
# check trailing dimensions
@@ -167,8 +185,17 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
sensordata = np.empty((nroll, nstep, nsensordata))
# call rollout
_rollout.rollout(model, data, nstep, control_spec, initial_state,
initial_warmstart, control, state, sensordata)
_rollout.rollout(
model,
data,
nstep,
control_spec,
initial_state,
initial_warmstart,
control,
state,
sensordata,
)
# return outputs
return state, sensordata
@@ -227,8 +254,8 @@ def _infer_dimension(dim, value, **kwargs):
Args:
dim: Dimension to be inferred.
value: Initial guess of inferred value (1: unknown).
**kwargs: List of arrays which should all have the same size (or 1)
along dimension dim.
**kwargs: List of arrays which should all have the same size (or 1) along
dimension dim.
Returns:
Inferred dimension.
python/mujoco/rollout_test.py
+125 -74
View File
@@ -127,10 +127,12 @@ TEST_XML_DIVERGE = r"""
</mujoco>
"""
ALL_MODELS = {'TEST_XML': TEST_XML,
'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
'TEST_XML_EMPTY': TEST_XML_EMPTY}
ALL_MODELS = {
'TEST_XML': TEST_XML,
'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
'TEST_XML_EMPTY': TEST_XML_EMPTY,
}
# ------------------------------ tests -----------------------------------------
@@ -242,8 +244,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.random.randn(nstate)
control = np.random.randn(nstep, model.nu)
initial_warmstart = np.tile(data.qacc_warmstart.copy(), (nroll, 1))
state, sensordata = rollout.rollout(model, data, initial_state, control,
initial_warmstart=initial_warmstart)
state, sensordata = rollout.rollout(
model, data, initial_state, control, initial_warmstart=initial_warmstart
)
mujoco.mj_resetData(model, data)
initial_state = np.tile(initial_state, (nroll, 1))
@@ -264,8 +267,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.random.randn(nstate)
control = np.random.randn(nstep, model.nu)
state = np.empty((nroll, nstep, nstate))
state, sensordata = rollout.rollout(model, data, initial_state, control,
state=state)
state, sensordata = rollout.rollout(
model, data, initial_state, control, state=state
)
mujoco.mj_resetData(model, data)
initial_state = np.tile(initial_state, (nroll, 1))
@@ -286,8 +290,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.random.randn(nstate)
control = np.random.randn(nstep, model.nu)
sensordata = np.empty((nroll, nstep, model.nsensordata))
state, sensordata = rollout.rollout(model, data, initial_state, control,
sensordata=sensordata)
state, sensordata = rollout.rollout(
model, data, initial_state, control, sensordata=sensordata
)
mujoco.mj_resetData(model, data)
initial_state = np.tile(initial_state, (nroll, 1))
@@ -309,8 +314,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
control = np.random.randn(model.nu)
state = np.empty((nroll, nstep, nstate))
sensordata = np.empty((nroll, nstep, model.nsensordata))
rollout.rollout(model, data, initial_state, control,
state=state, sensordata=sensordata)
rollout.rollout(
model, data, initial_state, control, state=state, sensordata=sensordata
)
control = np.tile(control, (nstep, 1))
py_state, py_sensordata = py_rollout(model, data, initial_state, control)
@@ -374,8 +380,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(nroll, 1, model.nu)
state = np.empty((nroll, nstep, nstate))
state, sensordata = rollout.rollout(model, data, initial_state, control,
state=state)
state, sensordata = rollout.rollout(
model, data, initial_state, control, state=state
)
control = np.repeat(control, nstep, axis=1)
py_state, py_sensordata = py_rollout(model, data, initial_state, control)
@@ -393,17 +400,21 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.random.randn(nroll, nstate)
control_spec = (mujoco.mjtState.mjSTATE_CTRL |
mujoco.mjtState.mjSTATE_QFRC_APPLIED |
mujoco.mjtState.mjSTATE_XFRC_APPLIED)
control_spec = (
mujoco.mjtState.mjSTATE_CTRL
| mujoco.mjtState.mjSTATE_QFRC_APPLIED
| mujoco.mjtState.mjSTATE_XFRC_APPLIED
)
ncontrol = mujoco.mj_stateSize(model, control_spec)
control = np.random.randn(nroll, nstep, ncontrol)
state, sensordata = rollout.rollout(model, data, initial_state, control,
control_spec=control_spec)
state, sensordata = rollout.rollout(
model, data, initial_state, control, control_spec=control_spec
)
py_state, py_sensordata = py_rollout(model, data, initial_state, control,
control_spec=control_spec)
py_state, py_sensordata = py_rollout(
model, data, initial_state, control, control_spec=control_spec
)
np.testing.assert_array_equal(state, py_state)
np.testing.assert_array_equal(sensordata, py_sensordata)
@@ -416,15 +427,19 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.empty((nroll, nstate))
# get diverging (0, 2) and non-diverging (1, 3) states
mujoco.mj_getState(model, data, initial_state[0],
mujoco.mjtState.mjSTATE_FULLPHYSICS)
mujoco.mj_getState(model, data, initial_state[2],
mujoco.mjtState.mjSTATE_FULLPHYSICS)
mujoco.mj_getState(
model, data, initial_state[0], mujoco.mjtState.mjSTATE_FULLPHYSICS
)
mujoco.mj_getState(
model, data, initial_state[2], mujoco.mjtState.mjSTATE_FULLPHYSICS
)
mujoco.mj_resetDataKeyframe(model, data, 0) # keyframe 0 does not diverge
mujoco.mj_getState(model, data, initial_state[1],
mujoco.mjtState.mjSTATE_FULLPHYSICS)
mujoco.mj_getState(model, data, initial_state[3],
mujoco.mjtState.mjSTATE_FULLPHYSICS)
mujoco.mj_getState(
model, data, initial_state[1], mujoco.mjtState.mjSTATE_FULLPHYSICS
)
mujoco.mj_getState(
model, data, initial_state[3], mujoco.mjtState.mjSTATE_FULLPHYSICS
)
nstep = 10000 # divergence after ~15s, timestep = 2e-3
@@ -459,27 +474,40 @@ class MuJoCoRolloutTest(parameterized.TestCase):
thread_local.data = mujoco.MjData(model)
model_list = [model] * nroll
def call_rollout(initial_state, control, state, sensordata):
rollout.rollout(model_list, thread_local.data, initial_state, control,
skip_checks=True,
nstep=nstep, state=state, sensordata=sensordata)
rollout.rollout(
model_list,
thread_local.data,
initial_state,
control,
skip_checks=True,
nstep=nstep,
state=state,
sensordata=sensordata,
)
n = nroll // num_workers # integer division
chunks = [] # a list of tuples, one per worker
for i in range(num_workers-1):
chunks.append((initial_state[i*n:(i+1)*n],
control[i*n:(i+1)*n],
state[i*n:(i+1)*n],
sensordata[i*n:(i+1)*n]))
for i in range(num_workers - 1):
chunks.append((
initial_state[i * n : (i + 1) * n],
control[i * n : (i + 1) * n],
state[i * n : (i + 1) * n],
sensordata[i * n : (i + 1) * n],
))
# last chunk, absorbing the remainder:
chunks.append((initial_state[(num_workers-1)*n:],
control[(num_workers-1)*n:],
state[(num_workers-1)*n:],
sensordata[(num_workers-1)*n:]))
chunks.append((
initial_state[(num_workers - 1) * n :],
control[(num_workers - 1) * n :],
state[(num_workers - 1) * n :],
sensordata[(num_workers - 1) * n :],
))
with concurrent.futures.ThreadPoolExecutor(
max_workers=num_workers, initializer=thread_initializer) as executor:
max_workers=num_workers, initializer=thread_initializer
) as executor:
futures = []
for chunk in chunks:
futures.append(executor.submit(call_rollout, *chunk))
@@ -513,12 +541,14 @@ class MuJoCoRolloutTest(parameterized.TestCase):
state, _ = rollout.rollout(model, data, state1[0], control)
# assert that stepping without warmstarts is not exact
np.testing.assert_raises(AssertionError,
np.testing.assert_array_equal, state, state2)
np.testing.assert_raises(
AssertionError, np.testing.assert_array_equal, state, state2
)
# take step using rollout, take warmstart into account
state, _ = rollout.rollout(model, data, state1, control,
initial_warmstart=initial_warmstart)
state, _ = rollout.rollout(
model, data, state1, control, initial_warmstart=initial_warmstart
)
# assert exact equality
np.testing.assert_array_equal(state, np.expand_dims(state2, axis=0))
@@ -530,19 +560,21 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.zeros(nstate)
control_spec = (mujoco.mjtState.mjSTATE_MOCAP_POS |
mujoco.mjtState.mjSTATE_MOCAP_QUAT)
control_spec = (
mujoco.mjtState.mjSTATE_MOCAP_POS | mujoco.mjtState.mjSTATE_MOCAP_QUAT
)
pos1 = np.array((1., 2., 3.))
quat1 = np.array((1., 2., 3., 4.))
pos1 = np.array((1.0, 2.0, 3.0))
quat1 = np.array((1.0, 2.0, 3.0, 4.0))
quat1 /= np.linalg.norm(quat1)
pos2 = np.array((2., 3., 4.))
quat2 = np.array((2., 3., 4., 5.))
pos2 = np.array((2.0, 3.0, 4.0))
quat2 = np.array((2.0, 3.0, 4.0, 5.0))
quat2 /= np.linalg.norm(quat2)
control = np.hstack((pos1, pos2, quat1, quat2))
_, sensordata = rollout.rollout(model, data, initial_state, control,
control_spec=control_spec)
_, sensordata = rollout.rollout(
model, data, initial_state, control, control_spec=control_spec
)
np.testing.assert_array_almost_equal(sensordata[0][0][:3], pos1)
np.testing.assert_array_almost_equal(sensordata[0][0][3:], quat1)
@@ -562,7 +594,8 @@ class MuJoCoRolloutTest(parameterized.TestCase):
model.opt.solver = 10 # invalid solver type
with self.assertRaisesWithLiteralMatch(
mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'):
mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'
):
rollout.rollout(model, data, initial_state, ctrl)
def test_invalid(self):
@@ -576,12 +609,14 @@ class MuJoCoRolloutTest(parameterized.TestCase):
control = 'string'
with self.assertRaisesWithLiteralMatch(
ValueError, 'control must be a numpy array or float'):
ValueError, 'control must be a numpy array or float'
):
rollout.rollout(model, data, initial_state, control)
control = np.zeros((2, 3, 4, 5))
with self.assertRaisesWithLiteralMatch(
ValueError, 'control can have at most 3 dimensions'):
ValueError, 'control can have at most 3 dimensions'
):
rollout.rollout(model, data, initial_state, control)
def test_bad_sizes(self):
@@ -594,28 +629,33 @@ class MuJoCoRolloutTest(parameterized.TestCase):
initial_state = np.random.randn(nroll, nstate + 1)
with self.assertRaisesWithLiteralMatch(
ValueError, 'trailing dimension of initial_state must be 6, got 7'):
ValueError, 'trailing dimension of initial_state must be 6, got 7'
):
rollout.rollout(model, data, initial_state)
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(1, nstep, model.nu + 1)
with self.assertRaisesWithLiteralMatch(
ValueError, 'trailing dimension of control must be 2, got 3'):
ValueError, 'trailing dimension of control must be 2, got 3'
):
rollout.rollout(model, data, initial_state, control)
control = np.random.randn(nroll, nstep, model.nu)
state = np.random.randn(nroll, nstep+1, nstate) # incompatible nstep
state = np.random.randn(nroll, nstep + 1, nstate) # incompatible nstep
with self.assertRaisesWithLiteralMatch(
ValueError, 'dimension 1 inferred as 3 but state has 4'):
ValueError, 'dimension 1 inferred as 3 but state has 4'
):
rollout.rollout(model, data, initial_state, control, state=state)
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(nroll, nstep, model.nu)
bad_spec = mujoco.mjtState.mjSTATE_ACT
with self.assertRaisesWithLiteralMatch(
ValueError, 'control_spec can only contain bits in mjSTATE_USER'):
rollout.rollout(model, data, initial_state, control,
control_spec=bad_spec)
ValueError, 'control_spec can only contain bits in mjSTATE_USER'
):
rollout.rollout(
model, data, initial_state, control, control_spec=bad_spec
)
def test_stateless(self):
model = mujoco.MjModel.from_xml_string(TEST_XML)
@@ -655,8 +695,9 @@ def get_state(model, data):
return state.reshape((1, nstate))
def step(model, data, state, control,
control_spec=mujoco.mjtState.mjSTATE_CTRL):
def step(
model, data, state, control, control_spec=mujoco.mjtState.mjSTATE_CTRL
):
if state is not None:
mujoco.mj_setState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS)
mujoco.mj_setState(model, data, control, control_spec)
@@ -664,8 +705,13 @@ def step(model, data, state, control,
return (get_state(model, data), data.sensordata)
def one_rollout(model, data, initial_state, control,
control_spec=mujoco.mjtState.mjSTATE_CTRL):
def one_rollout(
model,
data,
initial_state,
control,
control_spec=mujoco.mjtState.mjSTATE_CTRL,
):
nstep = control.shape[0]
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
state = np.empty((nstep, nstate))
@@ -673,9 +719,9 @@ def one_rollout(model, data, initial_state, control,
mujoco.mj_resetData(model, data)
for t in range(nstep):
state[t], sensordata[t] = step(model, data,
initial_state if t == 0 else None,
control[t], control_spec)
state[t], sensordata[t] = step(
model, data, initial_state if t == 0 else None, control[t], control_spec
)
return state, sensordata
@@ -700,15 +746,20 @@ def ensure_3d(arg):
return np.ascontiguousarray(arg, dtype=np.float64)
def py_rollout(model, data, initial_state, control,
control_spec=mujoco.mjtState.mjSTATE_CTRL):
def py_rollout(
model,
data,
initial_state,
control,
control_spec=mujoco.mjtState.mjSTATE_CTRL,
):
initial_state = ensure_2d(initial_state)
control = ensure_3d(control)
nroll = initial_state.shape[0]
nstep = control.shape[1]
if isinstance(model, mujoco.MjModel):
model = [model]*nroll
model = [model] * nroll
nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS)
python/mujoco/specs_test.py
+24 -11
View File
@@ -104,7 +104,9 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(model.nuser_site, 6)
np.testing.assert_array_equal(model.site_user[0], [1, 2, 3, 4, 5, 6])
self.assertEqual(spec.to_xml(), textwrap.dedent("""\
self.assertEqual(
spec.to_xml(),
textwrap.dedent("""\
<mujoco model="MuJoCo Model">
<compiler angle="radian"/>
@@ -116,7 +118,8 @@ class SpecsTest(absltest.TestCase):
</body>
</worldbody>
</mujoco>
"""),)
"""),
)
def test_kwarg(self):
# Create a spec.
@@ -467,7 +470,7 @@ class SpecsTest(absltest.TestCase):
# Try to compile, get error.
expected_error = (
'Error: size 0 must be positive in geom\n'
+ f'Element name \'MyGeom\', id 0, geom added on line {added_on_line}'
+ f"Element name 'MyGeom', id 0, geom added on line {added_on_line}"
)
with self.assertRaisesRegex(ValueError, expected_error):
spec.compile()
@@ -531,7 +534,9 @@ class SpecsTest(absltest.TestCase):
spec.worldbody.add_geom(main)
spec.compile()
self.assertEqual(spec.to_xml(), textwrap.dedent("""\
self.assertEqual(
spec.to_xml(),
textwrap.dedent("""\
<mujoco model="test">
<compiler angle="radian"/>
@@ -547,7 +552,8 @@ class SpecsTest(absltest.TestCase):
<geom/>
</worldbody>
</mujoco>
"""))
"""),
)
spec = mujoco.MjSpec()
spec.modelname = 'test'
@@ -561,7 +567,9 @@ class SpecsTest(absltest.TestCase):
spec.worldbody.add_geom(main)
spec.compile()
self.assertEqual(spec.to_xml(), textwrap.dedent("""\
self.assertEqual(
spec.to_xml(),
textwrap.dedent("""\
<mujoco model="test">
<compiler angle="radian"/>
@@ -577,7 +585,8 @@ class SpecsTest(absltest.TestCase):
<geom/>
</worldbody>
</mujoco>
"""))
"""),
)
def test_element_list(self):
spec = mujoco.MjSpec()
@@ -718,13 +727,17 @@ class SpecsTest(absltest.TestCase):
</worldbody>
</mujoco>
"""
spec = mujoco.MjSpec.from_string(textwrap.dedent("""
spec = mujoco.MjSpec.from_string(
textwrap.dedent("""
<mujoco model="MuJoCo Model">
<include file="included.xml"/>
</mujoco>
"""), {'included.xml': included_xml.encode('utf-8')})
self.assertEqual(spec.worldbody.first_body().first_geom().type,
mujoco.mjtGeom.mjGEOM_BOX)
"""),
{'included.xml': included_xml.encode('utf-8')},
)
self.assertEqual(
spec.worldbody.first_body().first_geom().type, mujoco.mjtGeom.mjGEOM_BOX
)
def test_delete(self):
file_path = epath.resource_path("mujoco") / "testdata" / "model.xml"
python/mujoco/viewer.py
+35 -22
View File
@@ -42,7 +42,7 @@ PERCENT_REALTIME = (
10, 8, 6.6, 5, 4, 3.3, 2.5, 2, 1.6, 1.3,
1, 0.8, 0.66, 0.5, 0.4, 0.33, 0.25, 0.2, 0.16, 0.13,
0.1
)
) # fmt: skip
# Maximum time mis-alignment before re-sync.
MAX_SYNC_MISALIGN = 0.1
@@ -194,12 +194,13 @@ def _file_loader(path: str) -> _LoaderWithPathType:
def _reload(
simulate: _Simulate, loader: _InternalLoaderType,
notify_loaded: Optional[Callable[[], None]] = None
simulate: _Simulate,
loader: _InternalLoaderType,
notify_loaded: Optional[Callable[[], None]] = None,
) -> Optional[Tuple[mujoco.MjModel, mujoco.MjData]]:
"""Internal function for reloading a model in the viewer."""
try:
simulate.load_message('') # path is unknown at this point
simulate.load_message('') # path is unknown at this point
load_tuple = loader()
except Exception as e: # pylint: disable=broad-except
simulate.load_error = str(e)
@@ -275,14 +276,16 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
# Inject noise.
if simulate.ctrl_noise_std != 0.0:
# Convert rate and scale to discrete time (Ornstein–Uhlenbeck).
rate = math.exp(-m.opt.timestep /
max(simulate.ctrl_noise_rate, mujoco.mjMINVAL))
rate = math.exp(
-m.opt.timestep / max(simulate.ctrl_noise_rate, mujoco.mjMINVAL)
)
scale = simulate.ctrl_noise_std * math.sqrt(1 - rate * rate)
for i in range(m.nu):
# Update noise.
ctrl_noise[i] = (rate * ctrl_noise[i] +
scale * mujoco.mju_standardNormal(None))
ctrl_noise[i] = rate * ctrl_noise[
i
] + scale * mujoco.mju_standardNormal(None)
# Apply noise.
d.ctrl[i] = ctrl_noise[i]
@@ -291,12 +294,18 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
slowdown = 100 / PERCENT_REALTIME[simulate.real_time_index]
# Misalignment: distance from target sim time > MAX_SYNC_MISALIGN.
misaligned = abs(elapsedcpu / slowdown -
elapsedsim) > MAX_SYNC_MISALIGN
misaligned = (
abs(elapsedcpu / slowdown - elapsedsim) > MAX_SYNC_MISALIGN
)
# Out-of-sync (for any reason): reset sync times, step.
if (elapsedsim < 0 or elapsedcpu < 0 or synccpu == 0 or misaligned or
simulate.speed_changed):
if (
elapsedsim < 0
or elapsedcpu < 0
or synccpu == 0
or misaligned
or simulate.speed_changed
):
# Re-sync.
synccpu = startcpu
syncsim = d.time
@@ -312,9 +321,9 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
prevsim = d.time
refreshtime = SIM_REFRESH_FRACTION / simulate.refresh_rate
# Step while sim lags behind CPU and within refreshtime.
while (((d.time - syncsim) * slowdown <
(time.time() - synccpu)) and
((time.time() - startcpu) < refreshtime)):
while (
(d.time - syncsim) * slowdown < (time.time() - synccpu)
) and ((time.time() - startcpu) < refreshtime):
# Measure slowdown before first step.
if not measured and elapsedsim:
simulate.measured_slowdown = elapsedcpu / elapsedsim
@@ -329,7 +338,7 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
break
# save current state to history buffer
if (stepped):
if stepped:
simulate.add_to_history()
else: # simulate.run is False: GUI is paused.
@@ -355,7 +364,8 @@ def _launch_internal(
raise ValueError('mjData is specified but mjModel is not')
elif callable(model) and data is not None:
raise ValueError(
'mjData should not be specified when an mjModel loader is used')
'mjData should not be specified when an mjModel loader is used'
)
elif loader is not None and model is not None:
raise ValueError('model and loader are both specified')
elif run_physics_thread and handle_return is not None:
@@ -398,14 +408,17 @@ def _launch_internal(
if run_physics_thread:
side_thread = threading.Thread(
target=_physics_loop, args=(simulate, loader))
target=_physics_loop, args=(simulate, loader)
)
else:
side_thread = threading.Thread(
target=_reload, args=(simulate, loader, notify_loaded))
target=_reload, args=(simulate, loader, notify_loaded)
)
def make_exit(simulate):
def exit_simulate():
simulate.exit()
return exit_simulate
exit_simulate = make_exit(simulate)
@@ -456,8 +469,7 @@ def launch_passive(
if not isinstance(data, mujoco.MjData):
raise ValueError(f'`data` is not a mujoco.MjData: got {data!r}')
if key_callback is not None and not callable(key_callback):
raise ValueError(
f'`key_callback` is not callable: got {key_callback!r}')
raise ValueError(f'`key_callback` is not callable: got {key_callback!r}')
mujoco.mj_forward(model, data)
handle_return = queue.Queue(1)
@@ -480,7 +492,8 @@ def launch_passive(
if not isinstance(_MJPYTHON, _MjPythonBase):
raise RuntimeError(
'`launch_passive` requires that the Python script be run under '
'`mjpython` on macOS')
'`mjpython` on macOS'
)
_MJPYTHON.launch_on_ui_thread(
model,
data,
python/pyproject.toml
+21
View File
@@ -65,3 +65,24 @@ usd = [
"usd-core",
"pillow"
]
[tool.isort]
force_single_line = true
force_sort_within_sections = true
lexicographical = true
single_line_exclusions = ["typing"]
order_by_type = false
group_by_package = true
line_length = 120
use_parentheses = true
multi_line_output = 3
skip_glob = ["**/*.ipynb"]
[tool.pyink]
line-length = 80
unstable = true
pyink-indentation = 2
pyink-use-majority-quotes = true
extend-exclude = '''(
.ipynb$
)'''
python/setup.py
+63 -40
View File
@@ -101,15 +101,15 @@ def tokenize_quoted_substr(input_string, quote_char, placeholders=None):
placeholders = placeholders if placeholders is not None else dict()
prev_end = -1
for start, end in start_and_end(quote_positions):
output_string += input_string[prev_end+1:start]
output_string += input_string[prev_end + 1 : start]
while True:
placeholder = ''.join(random.choices(string.ascii_lowercase, k=5))
if placeholder not in input_string and placeholder not in output_string:
break
output_string += placeholder
placeholders[placeholder] = input_string[start+1:end]
placeholders[placeholder] = input_string[start + 1 : end]
prev_end = end
output_string += input_string[prev_end+1:]
output_string += input_string[prev_end + 1 :]
return output_string, placeholders
@@ -145,15 +145,17 @@ class BuildCMakeExtension(build_ext.build_ext):
"""Uses CMake to build extensions."""
def run(self):
self._is_apple = (platform.system() == 'Darwin')
(self._mujoco_library_path,
self._mujoco_include_path,
self._mujoco_plugins_path,
self._mujoco_framework_path) = self._find_mujoco()
self._is_apple = platform.system() == 'Darwin'
(
self._mujoco_library_path,
self._mujoco_include_path,
self._mujoco_plugins_path,
self._mujoco_framework_path,
) = self._find_mujoco()
self._configure_cmake()
for ext in self.extensions:
assert ext.name.startswith(EXT_PREFIX)
assert '.' not in ext.name[len(EXT_PREFIX):]
assert '.' not in ext.name[len(EXT_PREFIX) :]
self.build_extension(ext)
self._copy_external_libraries()
self._copy_mujoco_headers()
@@ -163,20 +165,22 @@ class BuildCMakeExtension(build_ext.build_ext):
def _find_mujoco(self):
if MUJOCO_PATH not in os.environ:
raise RuntimeError(
f'{MUJOCO_PATH} environment variable is not set')
raise RuntimeError(f'{MUJOCO_PATH} environment variable is not set')
if MUJOCO_PLUGIN_PATH not in os.environ:
raise RuntimeError(
f'{MUJOCO_PLUGIN_PATH} environment variable is not set')
f'{MUJOCO_PLUGIN_PATH} environment variable is not set'
)
library_path = None
include_path = None
plugin_path = os.environ[MUJOCO_PLUGIN_PATH]
for directory, subdirs, filenames in os.walk(os.environ[MUJOCO_PATH]):
if self._is_apple and 'mujoco.framework' in subdirs:
return (os.path.join(directory, 'mujoco.framework/Versions/A'),
os.path.join(directory, 'mujoco.framework/Headers'),
plugin_path,
directory)
return (
os.path.join(directory, 'mujoco.framework/Versions/A'),
os.path.join(directory, 'mujoco.framework/Headers'),
plugin_path,
directory,
)
if fnmatch.filter(filenames, get_mujoco_lib_pattern()):
library_path = directory
if os.path.exists(os.path.join(directory, 'mujoco/mujoco.h')):
@@ -190,63 +194,78 @@ class BuildCMakeExtension(build_ext.build_ext):
for directory, _, filenames in os.walk(os.environ[MUJOCO_PATH]):
for pattern in get_external_lib_patterns():
for filename in fnmatch.filter(filenames, pattern):
shutil.copyfile(os.path.join(directory, filename),
os.path.join(dst, filename))
shutil.copyfile(
os.path.join(directory, filename), os.path.join(dst, filename)
)
def _copy_plugin_libraries(self):
dst = os.path.join(
os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
'plugin')
'plugin',
)
os.makedirs(dst)
for directory, _, filenames in os.walk(self._mujoco_plugins_path):
for pattern in get_plugin_lib_patterns():
for filename in fnmatch.filter(filenames, pattern):
shutil.copyfile(os.path.join(directory, filename),
os.path.join(dst, filename))
shutil.copyfile(
os.path.join(directory, filename), os.path.join(dst, filename)
)
def _copy_mujoco_headers(self):
dst = os.path.join(
os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
'include/mujoco')
'include/mujoco',
)
os.makedirs(dst)
for directory, _, filenames in os.walk(self._mujoco_include_path):
for filename in fnmatch.filter(filenames, '*.h'):
shutil.copyfile(os.path.join(directory, filename),
os.path.join(dst, filename))
shutil.copyfile(
os.path.join(directory, filename), os.path.join(dst, filename)
)
def _copy_mjpython(self):
src_dir = os.path.join(os.path.dirname(__file__), 'mujoco/mjpython')
dst_contents_dir = os.path.join(
os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
'MuJoCo_(mjpython).app/Contents')
'MuJoCo_(mjpython).app/Contents',
)
os.makedirs(dst_contents_dir)
shutil.copyfile(os.path.join(src_dir, 'Info.plist'),
os.path.join(dst_contents_dir, 'Info.plist'))
shutil.copyfile(
os.path.join(src_dir, 'Info.plist'),
os.path.join(dst_contents_dir, 'Info.plist'),
)
dst_bin_dir = os.path.join(dst_contents_dir, 'MacOS')
os.makedirs(dst_bin_dir)
shutil.copyfile(os.path.join(self.build_temp, 'mjpython'),
os.path.join(dst_bin_dir, 'mjpython'))
shutil.copyfile(
os.path.join(self.build_temp, 'mjpython'),
os.path.join(dst_bin_dir, 'mjpython'),
)
os.chmod(os.path.join(dst_bin_dir, 'mjpython'), 0o755)
dst_resources_dir = os.path.join(dst_contents_dir, 'Resources')
os.makedirs(dst_resources_dir)
shutil.copyfile(os.path.join(src_dir, 'mjpython.icns'),
os.path.join(dst_resources_dir, 'mjpython.icns'))
shutil.copyfile(
os.path.join(src_dir, 'mjpython.icns'),
os.path.join(dst_resources_dir, 'mjpython.icns'),
)
def _configure_cmake(self):
"""Check for CMake."""
cmake = os.environ.get(MUJOCO_CMAKE, 'cmake')
build_cfg = 'Debug' if self.debug else 'Release'
cmake_module_path = os.path.join(
os.path.dirname(__file__), 'mujoco', 'cmake')
os.path.dirname(__file__), 'mujoco', 'cmake'
)
cmake_args = [
f'-DPython3_ROOT_DIR:PATH={sys.prefix}',
f'-DPython3_EXECUTABLE:STRING={sys.executable}',
f'-DCMAKE_MODULE_PATH:PATH={cmake_module_path}',
f'-DCMAKE_BUILD_TYPE:STRING={build_cfg}',
f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY:PATH={self.build_temp}',
f'-DCMAKE_INTERPROCEDURAL_OPTIMIZATION:BOOL={"OFF" if self.debug else "ON"}',
(
f'-DCMAKE_INTERPROCEDURAL_OPTIMIZATION:BOOL={"OFF" if self.debug else "ON"}'
),
'-DCMAKE_Fortran_COMPILER:STRING=',
'-DBUILD_TESTING:BOOL=OFF',
]
@@ -284,14 +303,17 @@ class BuildCMakeExtension(build_ext.build_ext):
for arg in cmake_args:
print(f' {arg}')
subprocess.check_call(
[cmake] + cmake_args +
[os.path.join(os.path.dirname(__file__), 'mujoco')],
cwd=self.build_temp)
[cmake]
+ cmake_args
+ [os.path.join(os.path.dirname(__file__), 'mujoco')],
cwd=self.build_temp,
)
print('Building all extensions with CMake')
subprocess.check_call(
[cmake, '--build', '.', f'-j{os.cpu_count()}', '--config', build_cfg],
cwd=self.build_temp)
cwd=self.build_temp,
)
def build_extension(self, ext):
dest_path = self.get_ext_fullpath(ext.name)
@@ -331,6 +353,7 @@ class InstallScripts(install_scripts.install_scripts):
else:
self.outfiles.append(oldfile)
setuptools.setup(
long_description=get_long_description(),
long_description_content_type='text/markdown',
@@ -350,7 +373,7 @@ setuptools.setup(
CMakeExtension('mujoco._specs'),
CMakeExtension('mujoco._structs'),
],
scripts=[
'mujoco/mjpython/mjpython.py'
] if platform.system() == 'Darwin' else [],
scripts=['mujoco/mjpython/mjpython.py']
if platform.system() == 'Darwin'
else [],
)