Support float32 precision for mjtNum in Python bindings.

PiperOrigin-RevId: 922959317
Change-Id: I384b3bdeabdc4d3f66c71cae41dabfaf51387be7

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:

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"""
 <mujoco model="test">
   <compiler coordinate="local" angle="radian" eulerseq="xyz"/>
@@ -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"""

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':

python/mujoco/constants.cc

@@ -15,6 +15,7 @@
 #include <utility>
 #include <vector>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
 #include <pybind11/cast.h>
@@ -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

python/mujoco/serialization.h

@@ -17,6 +17,7 @@
 
 #include <cstddef>
 #include <iostream>
+#include <type_traits>
 
 #include <mujoco/mjtype.h>
 
@@ -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<mjtNum, float> && sizeof(mjtNum) == 4) ||
+              (std::is_same_v<mjtNum, double> && sizeof(mjtNum) == 8));
 
 inline void WriteChar(std::ostream& output, char c) {
   output.write(&c, 1);

python/mujoco/specs.cc

@@ -59,6 +59,9 @@ using MjDouble9 = Eigen::Map<Eigen::Matrix<double, 9, 1>>;
 using MjDouble10 = Eigen::Map<Eigen::Matrix<double, 10, 1>>;
 using MjDouble11 = Eigen::Map<Eigen::Matrix<double, 11, 1>>;
 using MjDoubleVec = Eigen::Map<Eigen::VectorXd>;
+using MjNum2 = Eigen::Map<Eigen::Matrix<mjtNum, 2, 1>>;
+using MjNum3 = Eigen::Map<Eigen::Matrix<mjtNum, 3, 1>>;
+using MjNum5 = Eigen::Map<Eigen::Matrix<mjtNum, 5, 1>>;
 
 using MjIntRef2 = Eigen::Ref<const Eigen::Vector2i>;
 using MjIntRef3 = Eigen::Ref<const Eigen::Vector3i>;
@@ -74,6 +77,9 @@ using MjDoubleRef9 = Eigen::Ref<const Eigen::Matrix<double, 9, 1>>;
 using MjDoubleRef10 = Eigen::Ref<const Eigen::Matrix<double, 10, 1>>;
 using MjDoubleRef11 = Eigen::Ref<const Eigen::Matrix<double, 11, 1>>;
 using MjDoubleRefVec = Eigen::Ref<const Eigen::VectorXd>;
+using MjNumRef2 = Eigen::Ref<const Eigen::Matrix<mjtNum, 2, 1>>;
+using MjNumRef3 = Eigen::Ref<const Eigen::Matrix<mjtNum, 3, 1>>;
+using MjNumRef5 = Eigen::Ref<const Eigen::Matrix<mjtNum, 5, 1>>;
 
 template <typename LoadFunc>
 static raw::MjSpec* LoadSpecFileImpl(