diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py index 799a13a5..fdd86372 100644 --- a/python/mujoco/__init__.py +++ b/python/mujoco/__init__.py @@ -59,6 +59,9 @@ from mujoco._functions import * from mujoco._specs import * from mujoco._structs import * +import numpy as np +MJTNUM_DTYPE = np.float32 if MJTNUM_BYTES == 4 else np.float64 # pylint: disable=undefined-variable + try: from mujoco._render import * # pytype: disable=import-error except ImportError: diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index a79b2d1a..cd1ab24f 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -25,6 +25,8 @@ from absl.testing import parameterized import mujoco import numpy as np +DTYPE = mujoco.MJTNUM_DTYPE + TEST_XML = r""" @@ -155,21 +157,21 @@ class MuJoCoBindingsTest(parameterized.TestCase): ) def test_can_read_array(self): - np.testing.assert_array_equal( + np.testing.assert_allclose( self.model.body_pos, [[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( + np.testing.assert_allclose( 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_allclose(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 @@ -340,24 +342,22 @@ class MuJoCoBindingsTest(parameterized.TestCase): ) def test_named_indexing_repr_in_data(self): - 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'))) + view_repr = repr(self.data.geom('mybox')) + self.assertStartsWith(view_repr, '<_MjDataGeomViews') + self.assertIn("name: 'mybox'", view_repr) + self.assertIn('xpos:', view_repr) + self.assertEndsWith(view_repr, '>') def test_named_indexing_body_repr_in_data(self): view_repr = repr(self.data.body('mybox')) self.assertStartsWith(view_repr, '<_MjDataBodyViews') - self.assertIn('xpos: array([0., 0., 0.])', view_repr) + self.assertIn('xpos:', view_repr) self.assertEndsWith(view_repr, '>') def test_named_indexing_repr_in_model(self): view_repr = repr(self.model.geom('mybox')) self.assertStartsWith(view_repr, '<_MjModelGeomViews') - self.assertIn('size: array([0.1, 0.1, 0.1])', view_repr) + self.assertIn('size:', view_repr) self.assertEndsWith(view_repr, '>') def test_addresses_differ_between_structs(self): @@ -385,25 +385,25 @@ class MuJoCoBindingsTest(parameterized.TestCase): self.assertIn('MjVisual', vis_repr) def test_mjmodel_can_read_and_write_opt(self): - self.assertEqual(self.model.opt.timestep, 0.002) - np.testing.assert_array_equal(self.model.opt.gravity, [0, 0, -9.81]) + np.testing.assert_allclose(self.model.opt.timestep, 0.002) + np.testing.assert_allclose(self.model.opt.gravity, [0, 0, -9.81]) opt = self.model.opt self.model.opt.timestep = 0.001 - self.assertEqual(opt.timestep, 0.001) + np.testing.assert_allclose(opt.timestep, 0.001) gravity = opt.gravity self.model.opt.gravity[1] = 0.1 - np.testing.assert_array_equal(gravity, [0, 0.1, -9.81]) + np.testing.assert_allclose(gravity, [0, 0.1, -9.81]) self.model.opt.gravity = 0.2 - np.testing.assert_array_equal(gravity, [0.2, 0.2, 0.2]) + np.testing.assert_allclose(gravity, [0.2, 0.2, 0.2]) def test_mjmodel_can_read_and_write_stat(self): self.assertNotEqual(self.model.stat.meanmass, 0) stat = self.model.stat self.model.stat.meanmass = 1.2 - self.assertEqual(stat.meanmass, 1.2) + np.testing.assert_allclose(stat.meanmass, 1.2) def test_mjmodel_can_read_and_write_vis(self): self.assertEqual(self.model.vis.quality.shadowsize, 51) @@ -435,8 +435,8 @@ class MuJoCoBindingsTest(parameterized.TestCase): def test_mjoption_can_make_default(self): opt = mujoco.MjOption() - self.assertEqual(opt.timestep, 0.002) - np.testing.assert_array_equal(opt.gravity, [0, 0, -9.81]) + np.testing.assert_allclose(opt.timestep, 0.002) + np.testing.assert_allclose(opt.gravity, [0, 0, -9.81]) def test_mjoption_can_copy(self): opt1 = mujoco.MjOption() @@ -444,14 +444,14 @@ class MuJoCoBindingsTest(parameterized.TestCase): opt1.gravity = 2 opt2 = copy.copy(opt1) - self.assertEqual(opt2.timestep, 0.001) - np.testing.assert_array_equal(opt2.gravity, [2, 2, 2]) + np.testing.assert_allclose(opt2.timestep, 0.001) + np.testing.assert_allclose(opt2.gravity, [2, 2, 2]) # Make sure opt2 is actually a copy. opt1.timestep = 0.005 opt1.gravity = 5 - self.assertEqual(opt2.timestep, 0.001) - np.testing.assert_array_equal(opt2.gravity, [2, 2, 2]) + np.testing.assert_allclose(opt2.timestep, 0.001) + np.testing.assert_allclose(opt2.gravity, [2, 2, 2]) def test_mjmodel_can_copy(self): model_copy = copy.copy(self.model) @@ -469,8 +469,8 @@ class MuJoCoBindingsTest(parameterized.TestCase): # Make sure it's a copy. self.model.geom_size[1] = 0.5 - np.testing.assert_array_equal(self.model.geom_size[1], [0.5, 0.5, 0.5]) - np.testing.assert_array_equal(model_copy.geom_size[1], [0.1, 0.1, 0.1]) + np.testing.assert_allclose(self.model.geom_size[1], [0.5, 0.5, 0.5]) + np.testing.assert_allclose(model_copy.geom_size[1], [0.1, 0.1, 0.1]) def test_mjdata_can_copy(self): self.data.qpos = [0, 0, 0.1*np.sqrt(2) - 0.001, @@ -598,21 +598,21 @@ class MuJoCoBindingsTest(parameterized.TestCase): expected_pos = [] for contact in self.data.contact: - expected_pos.append(np.random.uniform(size=3)) + expected_pos.append(np.array(np.random.uniform(size=3), dtype=DTYPE)) contact.pos = expected_pos[-1] self.assertLen(expected_pos, expected_ncon) np.testing.assert_array_equal(self.data.contact.pos, expected_pos) expected_friction = [] for contact in self.data.contact: - expected_friction.append(np.random.uniform(size=5)) + expected_friction.append(np.array(np.random.uniform(size=5), dtype=DTYPE)) contact.friction = expected_friction[-1] self.assertLen(expected_friction, expected_ncon) np.testing.assert_array_equal(self.data.contact.friction, expected_friction) expected_H = [] # pylint: disable=invalid-name for contact in self.data.contact: - expected_H.append(np.random.uniform(size=36)) + expected_H.append(np.array(np.random.uniform(size=36), dtype=DTYPE)) contact.H = expected_H[-1] self.assertLen(expected_H, expected_ncon) np.testing.assert_array_equal(self.data.contact.H, expected_H) @@ -749,7 +749,7 @@ class MuJoCoBindingsTest(parameterized.TestCase): expected = np.array([1, 1, 0]) / np.sqrt(2) # Check that the output argument works, and that the binding returns None. - res = np.empty(3, np.float64) + res = np.empty(3, DTYPE) self.assertIsNone(mujoco.mju_rotVecQuat(res, vec, quat)) np.testing.assert_allclose(res, expected) @@ -759,7 +759,7 @@ class MuJoCoBindingsTest(parameterized.TestCase): # Check that the res argument must have the right size. with self.assertRaises(TypeError): - mujoco.mju_rotVecQuat(np.zeros(4, np.float64), vec, quat) + mujoco.mju_rotVecQuat(np.zeros(4, DTYPE), vec, quat) # Check that the vec argument must have the right size. with self.assertRaises(TypeError): @@ -792,24 +792,24 @@ class MuJoCoBindingsTest(parameterized.TestCase): sig = mujoco.mjtState.mjSTATE_INTEGRATION size = mujoco.mj_stateSize(self.model, sig) - state_bad_size = np.empty(size + 1, np.float64) + state_bad_size = np.empty(size + 1, DTYPE) expected_message = 'state size should equal mj_stateSize(m, sig)' with self.assertRaisesWithLiteralMatch(TypeError, expected_message): mujoco.mj_getState(self.model, self.data, state_bad_size, sig) # Get initial state. - state0 = np.empty(size, np.float64) + state0 = np.empty(size, DTYPE) mujoco.mj_getState(self.model, self.data, state0, sig) # Step, get next state. mujoco.mj_step(self.model, self.data) - state1a = np.empty(size, np.float64) + state1a = np.empty(size, DTYPE) mujoco.mj_getState(self.model, self.data, state1a, sig) # Reset to initial state, step again, get state again. mujoco.mj_setState(self.model, self.data, state0, sig) mujoco.mj_step(self.model, self.data) - state1b = np.empty(size, np.float64) + state1b = np.empty(size, DTYPE) mujoco.mj_getState(self.model, self.data, state1b, sig) # Expect next states to be equal. @@ -818,7 +818,7 @@ class MuJoCoBindingsTest(parameterized.TestCase): # Test mj_copyState data2 = mujoco.MjData(self.model) mujoco.mj_copyState(self.model, self.data, data2, sig) - state1c = np.empty(size, np.float64) + state1c = np.empty(size, DTYPE) mujoco.mj_getState(self.model, data2, state1c, sig) np.testing.assert_array_equal(state1a, state1c) @@ -853,14 +853,14 @@ class MuJoCoBindingsTest(parameterized.TestCase): np.testing.assert_array_equal(act, self.data.act) def test_mj_angmomMat(self): # pylint: disable=invalid-name - self.data.qvel = np.ones(self.model.nv, np.float64) + self.data.qvel = np.ones(self.model.nv, DTYPE) mujoco.mj_forward(self.model, self.data) mujoco.mj_subtreeVel(self.model, self.data) - mat = np.empty((3, 10), np.float64) + mat = np.empty((3, 10), DTYPE) 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_allclose( + mat @ self.data.qvel, self.data.subtree_angmom[0, :], atol=1e-5 ) def test_mj_jacSite(self): # pylint: disable=invalid-name @@ -868,7 +868,7 @@ class MuJoCoBindingsTest(parameterized.TestCase): site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite') # Call mj_jacSite with only jacp. - jacp = np.empty((3, 10), np.float64) + jacp = np.empty((3, 10), DTYPE) mujoco.mj_jacSite(self.model, self.data, jacp, None, site_id) expected_jacp = np.array([ @@ -879,7 +879,7 @@ class MuJoCoBindingsTest(parameterized.TestCase): np.testing.assert_array_equal(jacp, expected_jacp) # Call mj_jacSite with only jacr. - jacr = np.empty((3, 10), np.float64) + jacr = np.empty((3, 10), DTYPE) mujoco.mj_jacSite(self.model, self.data, None, jacr, site_id) expected_jacr = np.array([ @@ -941,8 +941,8 @@ Euler integrator, semi-implicit in velocity. ) def test_float_constant(self): - self.assertEqual(mujoco.mjMAXVAL, 1e10) - self.assertEqual(mujoco.mjMINVAL, 1e-15) + np.testing.assert_allclose(mujoco.mjMAXVAL, 1e10, rtol=1e-6) + np.testing.assert_allclose(mujoco.mjMINVAL, 1e-15, rtol=1e-6) def test_string_constants(self): self.assertLen(mujoco.mjDISABLESTRING, mujoco.mjtDisableBit.mjNDISABLE) @@ -1306,7 +1306,7 @@ Euler integrator, semi-implicit in velocity. geomid, ) # Check that named arguments work - normal = np.zeros(3, np.float64) + normal = np.zeros(3, DTYPE) mujoco.mj_ray( m=self.model, d=self.data, @@ -1321,14 +1321,14 @@ Euler integrator, semi-implicit in velocity. def test_mju_ray_geom(self): # Test mju_rayGeom with a plane at origin - pos = np.zeros(3) - mat = np.eye(3).flatten() - size = np.array([10.0, 10.0, 1.0]) - pnt = np.array([5.0, 5.0, 5.0]) + pos = np.zeros(3, DTYPE) + mat = np.eye(3, dtype=DTYPE).flatten() + size = np.array([10.0, 10.0, 1.0], dtype=DTYPE) + pnt = np.array([5.0, 5.0, 5.0], dtype=DTYPE) # Normalize direction for Euclidean distance - vec = np.array([-1.0, -1.0, -1.0]) + vec = np.array([-1.0, -1.0, -1.0], dtype=DTYPE) vec = vec / np.linalg.norm(vec) - normal = np.zeros(3) + normal = np.zeros(3, DTYPE) # Call with normal argument dist = mujoco.mju_rayGeom( @@ -1347,12 +1347,12 @@ Euler integrator, semi-implicit in velocity. def test_mj_multi_ray(self): nray = 3 geom1 = np.zeros(1, np.int32) - pnt = np.array([-0.3, 0, 0.1]) - vec = np.array([[1, 0, 0], [0, 0, 1], [0, 0, -1]], np.float64) + pnt = np.array([-0.3, 0, 0.1], DTYPE) + vec = np.array([[1, 0, 0], [0, 0, 1], [0, 0, -1]], DTYPE) dist_ex = np.array([0.2, -1, 0.1]) geom_ex = np.array([1, -1, 0]) geomid = np.zeros(nray, np.int32) - dist = np.zeros(nray, np.float64) + dist = np.zeros(nray, DTYPE) mujoco.mj_forward(self.model, self.data) mujoco.mj_multiRay( @@ -1400,23 +1400,23 @@ Euler integrator, semi-implicit in velocity. def test_geom_distance(self): mujoco.mj_forward(self.model, self.data) - fromto = np.empty(6, np.float64) + fromto = np.empty(6, DTYPE) 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.0, 0.0, 42.0, 0.0, 41.9)) + np.testing.assert_allclose(dist, 41.9, atol=1e-5) + np.testing.assert_allclose( + fromto, np.array((42.0, 0.0, 0.0, 42.0, 0.0, 41.9)), atol=1e-5 ) def test_inverse_fd(self): eps = 1e-6 flg_centered = 0 - df_dq = np.zeros((self.model.nv, self.model.nv)) - df_dv = np.zeros((self.model.nv, self.model.nv)) - df_da = np.zeros((self.model.nv, self.model.nv)) - ds_dq = np.zeros((self.model.nv, self.model.nsensordata)) - 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)) + df_dq = np.zeros((self.model.nv, self.model.nv), dtype=DTYPE) + df_dv = np.zeros((self.model.nv, self.model.nv), dtype=DTYPE) + df_da = np.zeros((self.model.nv, self.model.nv), dtype=DTYPE) + ds_dq = np.zeros((self.model.nv, self.model.nsensordata), dtype=DTYPE) + ds_dv = np.zeros((self.model.nv, self.model.nsensordata), dtype=DTYPE) + ds_da = np.zeros((self.model.nv, self.model.nsensordata), dtype=DTYPE) + dm_dq = np.zeros((self.model.nv, self.model.nM), dtype=DTYPE) mujoco.mjd_inverseFD( self.model, self.data, @@ -1440,11 +1440,11 @@ Euler integrator, semi-implicit in velocity. def test_mjd_sub_quat(self): quat1 = np.array((0.2, 0.3, 0.3, 0.4)) quat2 = np.array((0.1, 0.2, 0.4, 0.5)) - d1 = np.empty(9, np.float64) - d2 = np.empty(9, np.float64) + d1 = np.empty(9, DTYPE) + d2 = np.empty(9, DTYPE) mujoco.mjd_subQuat(quat1, quat2, d1, d2) - d3 = np.empty((3, 3), np.float64) - d4 = np.empty((3, 3), np.float64) + d3 = np.empty((3, 3), DTYPE) + d4 = np.empty((3, 3), DTYPE) mujoco.mjd_subQuat(quat1, quat2, None, d3) mujoco.mjd_subQuat(quat1, quat2, d4, None) np.testing.assert_array_equal(d2, d3.flatten()) @@ -1453,9 +1453,9 @@ Euler integrator, semi-implicit in velocity. def test_mjd_quat_integrate(self): scale = 0.1 vel = np.array((0.2, 0.3, 0.3)) - d_quat = np.empty(9, np.float64) - d_vel = np.empty(9, np.float64) - d_h = np.empty(3, np.float64) + d_quat = np.empty(9, DTYPE) + d_vel = np.empty(9, DTYPE) + d_h = np.empty(3, DTYPE) mujoco.mjd_quatIntegrate(vel, scale, d_quat, d_vel, d_h) def test_banded(self): @@ -1467,23 +1467,23 @@ Euler integrator, semi-implicit in velocity. [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) + ], dtype=DTYPE) + band = np.zeros(n_band * (n_total - n_dense) + n_dense * n_total, dtype=DTYPE) 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) - dense2 = np.zeros((n_total, n_total)) + dense2 = np.zeros((n_total, n_total), dtype=DTYPE) flg_sym = 1 mujoco.mju_band2Dense(dense2, band, n_total, n_band, n_dense, flg_sym) np.testing.assert_array_equal(dense, dense2) - vec = np.array([[2.0], [2.0], [3.0], [4.0]]) + vec = np.array([[2.0], [2.0], [3.0], [4.0]], dtype=DTYPE) res = np.zeros_like(vec) n_vec = 1 mujoco.mju_bandMulMatVec( res, band, vec, n_total, n_band, n_dense, n_vec, flg_sym ) - np.testing.assert_array_equal(res, dense @ vec) + np.testing.assert_allclose(res, dense @ vec) diag_add = 0 diag_mul = 0 mujoco.mju_cholFactorBand( @@ -1494,31 +1494,31 @@ Euler integrator, semi-implicit in velocity. def test_mju_box_qp(self): n = 5 - res = np.zeros(n) - r = np.zeros((n, n + 7)) + res = np.zeros(n, DTYPE) + r = np.zeros((n, n + 7), DTYPE) index = np.zeros(n, np.int32) - h = np.eye(n) - g = np.ones((n,)) - lower = -np.ones((n,)) - upper = np.ones((n,)) + h = np.eye(n, dtype=DTYPE) + g = np.ones((n,), dtype=DTYPE) + lower = -np.ones((n,), dtype=DTYPE) + upper = np.ones((n,), dtype=DTYPE) rank = mujoco.mju_boxQP(res, r, index, h, g, lower, upper) self.assertGreater(rank, -1) def test_mju_fill(self): - res = np.empty(3, np.float64) + res = np.empty(3, DTYPE) mujoco.mju_fill(res, 1.5) - np.testing.assert_array_equal(res, np.full(3, 1.5)) + np.testing.assert_allclose(res, np.full(3, 1.5)) def test_mju_eye(self): - eye4 = np.empty((4, 4), np.float64) + eye4 = np.empty((4, 4), DTYPE) mujoco.mju_eye(eye4) np.testing.assert_array_equal(eye4, np.eye(4)) def test_mju_symmetrize(self): - mat = np.linspace(0, 1, 16).reshape(4, 4) - res = np.empty((4, 4), np.float64) + mat = np.linspace(0, 1, 16, dtype=DTYPE).reshape(4, 4) + res = np.empty((4, 4), DTYPE) mujoco.mju_symmetrize(res, mat) - np.testing.assert_array_equal(res, 0.5 * (mat + mat.T)) + np.testing.assert_allclose(res, 0.5 * (mat + mat.T)) def test_mju_clip(self): self.assertEqual(mujoco.mju_clip(1.5, 1.0, 2.0), 1.5) @@ -1532,12 +1532,12 @@ Euler integrator, semi-implicit in velocity. self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.0) def test_mju_dense_to_sparse(self): - mat = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]]) + mat = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]], dtype=DTYPE) 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]) - vals = np.zeros(3) + vals = np.zeros(3, DTYPE) row_nnz = np.zeros(2, np.int32) row_adr = np.zeros(2, np.int32) col_ind = np.zeros(3, np.int32) @@ -1550,21 +1550,21 @@ Euler integrator, semi-implicit in velocity. def test_mju_sparse_to_dense(self): expected = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]]) - mat = np.array((1.0, 2.0, 3.0)) + mat = np.array((1.0, 2.0, 3.0), dtype=DTYPE) rownnz = np.array([1, 2]) rowadr = np.array([0, 1]) colind = np.array([1, 0, 2]) - res = np.zeros((2, 3)) + res = np.zeros((2, 3), dtype=DTYPE) mujoco.mju_sparse2dense(res, mat, rownnz, rowadr, colind) np.testing.assert_array_equal(res, expected) def test_mju_euler_to_quat(self): - quat = np.zeros(4) - euler = np.array([0, np.pi / 2, 0]) + quat = np.zeros(4, DTYPE) + euler = np.array([0, np.pi / 2, 0], DTYPE) seq = 'xyz' mujoco.mju_euler2Quat(quat, euler, seq) expected_quat = np.array([np.sqrt(0.5), 0, np.sqrt(0.5), 0.0]) - np.testing.assert_almost_equal(quat, expected_quat) + np.testing.assert_almost_equal(quat, expected_quat, decimal=5) error = 'mju_euler2Quat: seq must contain exactly 3 characters' with self.assertRaisesWithLiteralMatch(mujoco.FatalError, error): @@ -1798,21 +1798,21 @@ Euler integrator, semi-implicit in velocity. self.assertEqual(nhistory, 3) self.assertEqual(dim, 3) # accelerometer has dim=3 times = np.array([0.0, 0.01, 0.02]) - values = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float64) + values = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=DTYPE) mujoco.mj_initSensorHistory(model, data, 0, times, values, phase=0.0) # Read back a value using zero-order hold # mj_readSensor auto-subtracts delay: lookup_time = read_time - delay # delay = 0.01, so: # read_time=0.02 -> lookup at 0.01 -> value [4, 5, 6] - result = np.zeros(dim) + result = np.zeros(dim, DTYPE) mujoco.mj_readSensor(model, data, 0, 0.02, result, interp=0) # ZOH returns value at t=0.01 np.testing.assert_array_equal(result, [4, 5, 6]) # Test with times=None (uses existing timestamps) new_values = np.array([ - [10, 11, 12], [13, 14, 15], [16, 17, 18]], dtype=np.float64) + [10, 11, 12], [13, 14, 15], [16, 17, 18]], dtype=DTYPE) mujoco.mj_initSensorHistory(model, data, 0, None, new_values, phase=0.0) # read_time=0.02 -> lookup at 0.01 -> value [13, 14, 15] mujoco.mj_readSensor(model, data, 0, 0.02, result, interp=0) @@ -1861,7 +1861,7 @@ Euler integrator, semi-implicit in velocity. # After mj_makeData, reading from the delay buffer gives all zeros. # For a quaternion sensor, this is invalid data. - result = np.zeros(dim) + result = np.zeros(dim, DTYPE) mujoco.mj_readSensor(model, data, 0, delay, result, interp=0) np.testing.assert_array_equal(result, [0, 0, 0, 0]) @@ -1883,7 +1883,7 @@ Euler integrator, semi-implicit in velocity. # Now reading from the delay buffer gives the valid identity quaternion. mujoco.mj_readSensor(model, data, 0, delay, result, interp=0) - np.testing.assert_array_equal(result, [1, 0, 0, 0]) + np.testing.assert_allclose(result, [1, 0, 0, 0]) def test_poly_stiffness(self): xml = r""" diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 5dee5684..2b1d4b14 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -156,8 +156,10 @@ def _array_binding_code( if classname == 'mjSpec': # raw mjSpec has a wrapper rawclassname = classname.replace('mjS', 'MjS') fullvarname = 'ptr->' + varname - if innertype == 'double' or innertype == 'mjtNum': - innertype = 'MjDouble' # custom Eigen type + if innertype == 'mjtNum': + innertype = 'MjNum' # custom Eigen type for mjtNum fields + elif innertype == 'double': + innertype = 'MjDouble' # custom Eigen type for double fields elif innertype == 'float': innertype = 'MjFloat' # custom Eigen type elif innertype == 'int': diff --git a/python/mujoco/constants.cc b/python/mujoco/constants.cc index 593070d2..023b78cb 100644 --- a/python/mujoco/constants.cc +++ b/python/mujoco/constants.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -94,6 +95,7 @@ PYBIND11_MODULE(_constants, pymodule) { pymodule.attr("mjFRAMESTRING") = MakeTuple(mjFRAMESTRING); pymodule.attr("mjVISSTRING") = MakeTuple(mjVISSTRING); pymodule.attr("mjRNDSTRING") = MakeTuple(mjRNDSTRING); + pymodule.attr("MJTNUM_BYTES") = sizeof(mjtNum); } } // namespace } // namespace mujoco::python diff --git a/python/mujoco/serialization.h b/python/mujoco/serialization.h index cb234ac8..b387b455 100644 --- a/python/mujoco/serialization.h +++ b/python/mujoco/serialization.h @@ -17,6 +17,7 @@ #include #include +#include #include @@ -30,7 +31,8 @@ namespace mujoco::python::_impl { static_assert(sizeof(char) == 1); static_assert(sizeof(int) == 4); -static_assert(sizeof(mjtNum) == 8); +static_assert((std::is_same_v && sizeof(mjtNum) == 4) || + (std::is_same_v && sizeof(mjtNum) == 8)); inline void WriteChar(std::ostream& output, char c) { output.write(&c, 1); diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc index 738ab73b..f8abf3ad 100644 --- a/python/mujoco/specs.cc +++ b/python/mujoco/specs.cc @@ -59,6 +59,9 @@ using MjDouble9 = Eigen::Map>; using MjDouble10 = Eigen::Map>; using MjDouble11 = Eigen::Map>; using MjDoubleVec = Eigen::Map; +using MjNum2 = Eigen::Map>; +using MjNum3 = Eigen::Map>; +using MjNum5 = Eigen::Map>; using MjIntRef2 = Eigen::Ref; using MjIntRef3 = Eigen::Ref; @@ -74,6 +77,9 @@ using MjDoubleRef9 = Eigen::Ref>; using MjDoubleRef10 = Eigen::Ref>; using MjDoubleRef11 = Eigen::Ref>; using MjDoubleRefVec = Eigen::Ref; +using MjNumRef2 = Eigen::Ref>; +using MjNumRef3 = Eigen::Ref>; +using MjNumRef5 = Eigen::Ref>; template static raw::MjSpec* LoadSpecFileImpl(