Allow for different MJX backend implementations.

PiperOrigin-RevId: 755935704
Change-Id: Ic135cd00137c2857c73c683ed9fdc5ac4418715d

mjx/mujoco/mjx/_src/collision_convex.py

@@ -28,8 +28,10 @@ from mujoco.mjx._src.collision_types import FunctionKey
 from mujoco.mjx._src.collision_types import GeomInfo
 from mujoco.mjx._src.collision_types import HFieldInfo
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import GeomType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 # pylint: enable=g-importing-member
 
 _GeomInfo = Union[GeomInfo, ConvexInfo]
@@ -42,6 +44,9 @@ def collider(ncon: int):
     def collide(
         m: Model, d: Data, key: FunctionKey, geom: jax.Array
     ) -> Collision:
+      if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+        raise ValueError('collider requires JAX backend implementation.')
+
       g1, g2 = geom.T
       infos = [
           GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]),
@@ -56,7 +61,7 @@ def collider(ncon: int):
               pos=0, mat=0, size=0, face=0, vert=0
           )
         elif key.types[i] == GeomType.MESH:
-          c, cm = infos[i], m.mesh_convex[key.data_ids[i]]
+          c, cm = infos[i], m._impl.mesh_convex[key.data_ids[i]]
           infos[i] = ConvexInfo(**vars(c), **vars(cm))
           in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
               pos=0, mat=0, size=0

mjx/mujoco/mjx/_src/collision_driver.py

@@ -71,9 +71,11 @@ from mujoco.mjx._src.collision_sdf import sphere_ellipsoid
 from mujoco.mjx._src.collision_types import FunctionKey
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import GeomType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -227,7 +229,7 @@ def _geom_groups(
     if types[0] == mujoco.mjtGeom.mjGEOM_HFIELD:
       # add static grid bounds to the grouping key for hfield collisions
       geom_rbound_hfield = (
-          m.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound
+          m._impl.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound  # pytype: disable=attribute-error
       )
       nrow, ncol = m.hfield_nrow[data_ids[0]], m.hfield_ncol[data_ids[0]]
       xsize, ysize = m.hfield_size[data_ids[0]][:2]
@@ -323,11 +325,11 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
         solref=solref,
         solreffriction=solreffriction,
         solimp=solimp,
-        dim=d.contact.dim,
+        dim=d._impl.contact.dim,  # pytype: disable=attribute-error
         geom1=jp.array(geom[:, 0]),
         geom2=jp.array(geom[:, 1]),
         geom=jp.array(geom[:, :2]),
-        efc_address=d.contact.efc_address,
+        efc_address=d._impl.contact.efc_address,  # pytype: disable=attribute-error
     )
 
   return groups
@@ -374,7 +376,10 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
 
 def collision(m: Model, d: Data) -> Data:
   """Collides geometries."""
-  if d.ncon == 0:
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('collision requires JAX backend implementation.')
+
+  if d._impl.ncon == 0:  # pytype: disable=attribute-error
     return d
 
   max_geom_pairs = _numeric(m, 'max_geom_pairs')
@@ -424,4 +429,4 @@ def collision(m: Model, d: Data) -> Data:
   contacts = sum([condim_groups[k] for k in sorted(condim_groups)], [])
   contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts)
 
-  return d.replace(contact=contact)
+  return d.replace(_impl=d._impl.replace(contact=contact))  # pytype: disable=attribute-error

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -119,7 +119,7 @@ class SphereCollisionTest(parameterized.TestCase):
   def test_sphere(self, name, mjcf):
     d, dx = _collide(mjcf)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-5)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-5)
 
   _SPHERE_CONVEX = """
     <mujoco>
@@ -143,7 +143,7 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     self.assertEmpty(d.contact.dist)
-    self.assertGreater(dx.contact.dist, 0)
+    self.assertGreater(dx._impl.contact.dist, 0)
 
     # face contact
     xml = self._SPHERE_CONVEX.replace(
@@ -151,38 +151,38 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'face', 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'face', 1e-4)
 
     # deep face contact
     xml = self._SPHERE_CONVEX.replace(
         '<body pos="0.52 0 0.52">', '<body pos="0.48 0 0.47">'
     )
     d, dx = _collide(xml)
-    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((dx._impl.contact.dist < 0).all())
     self.assertTrue((d.contact.dist < 0).all())
-    np.testing.assert_allclose(dx.contact.dist, [-0.07], atol=1e-5)
+    np.testing.assert_allclose(dx._impl.contact.dist, [-0.07], atol=1e-5)
-    np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
+    np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
+        dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3))
     )
 
   def test_sphere_convex_edge(self):
     # edge contact
     d, dx = _collide(self._SPHERE_CONVEX)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'edge', 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'edge', 1e-4)
 
     # deep edge penetration
     xml = self._SPHERE_CONVEX.replace(
         '<body pos="0.52 0 0.52">', '<body pos="0.49 0 0.49">'
     )
     d, dx = _collide(xml)
-    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((dx._impl.contact.dist < 0).all())
     self.assertTrue((d.contact.dist < 0).all())
-    np.testing.assert_allclose(dx.contact.dist, [-0.06], atol=1e-5)
+    np.testing.assert_allclose(dx._impl.contact.dist, [-0.06], atol=1e-5)
-    np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
+    np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
+        dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3))
     )
 
     # vertex contact
@@ -191,7 +191,7 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'vertex', 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'vertex', 1e-4)
 
     # sphere center on vertex
     xml = self._SPHERE_CONVEX.replace(
@@ -199,7 +199,9 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'vertex_center', 1e-4)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'vertex_center', 1e-4
+      )
 
 
 class EllipsoidCollisionTest(parameterized.TestCase):
@@ -219,10 +221,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
   def test_plane_ellipsoid(self):
     """Tests ellipsoid plane contact."""
     d, dx = _collide(self._ELLIPSOID_PLANE)
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
       )
 
   _ELLIPSOID_ELLIPSOID = """
@@ -242,10 +244,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
   def test_ellipsoid_ellipsoid(self):
     """Tests ellipsoid ellipsoid contact."""
     d, dx = _collide(self._ELLIPSOID_ELLIPSOID)
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2
       )
 
   _ELLIPSOID_SPHERE = """
@@ -266,10 +268,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     """Tests ellipsoid capsule contact."""
     d, dx = _collide(self._ELLIPSOID_SPHERE)
     d.contact.pos[0][2] = 0.03  # MJX finds the point on the surface
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4
       )
 
   _ELLIPSOID_CAPSULE = """
@@ -289,10 +291,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
   def test_capsule_ellipsoid(self):
     """Tests ellipsoid capsule contact."""
     d, dx = _collide(self._ELLIPSOID_CAPSULE)
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5
       )
 
   _ELLIPSOID_CYLINDER = """
@@ -313,10 +315,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     """Tests ellipsoid cylinder contact."""
     d, dx = _collide(self._ELLIPSOID_CYLINDER)
     d.contact.pos[0][2] = 0.04  # MJX finds the deepest point on the surface
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
       )
 
 
@@ -357,7 +359,7 @@ class CapsuleCollisionTest(parameterized.TestCase):
   def test_capsule(self, name, mjcf):
     d, dx = _collide(mjcf)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-4)
 
   _PARALLEL_CAP = """
     <mujoco>
@@ -378,14 +380,14 @@ class CapsuleCollisionTest(parameterized.TestCase):
     """Tests that two parallel capsules are colliding at the midpoint."""
     _, dx = _collide(self._PARALLEL_CAP)
 
-    np.testing.assert_allclose(dx.contact.dist, -0.05)
+    np.testing.assert_allclose(dx._impl.contact.dist, -0.05)
     np.testing.assert_allclose(
-        dx.contact.pos[0],
+        dx._impl.contact.pos[0],
         np.array([0.0, 0.1, (0.15 + 0.2) / 2.0]),
         atol=1e-5,
     )
     np.testing.assert_allclose(
-        dx.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5
+        dx._impl.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5
     )
 
   _CAP_BOX = """
@@ -408,30 +410,32 @@ class CapsuleCollisionTest(parameterized.TestCase):
     d, dx = _collide(self._CAP_BOX)
 
     # sort positions for comparison
-    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
-    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
     idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
     d.contact.pos[:] = d.contact.pos[idx]
     d.contact.frame[:] = d.contact.frame[idx]
     d.contact.dist[:] = d.contact.dist[idx]
 
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'capsule_convex', 1e-4
+      )
 
   def test_capsule_convex_face_deep(self):
     """Tests deep face penetration."""
     xml = self._CAP_BOX.replace('<body pos="0 0 0.54">', '<body pos="0 0 0.4">')
 
     _, dx = _collide(xml)
-    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((dx._impl.contact.dist < 0).all())
     np.testing.assert_array_almost_equal(
-        dx.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]])
+        dx._impl.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]])
     )
     np.testing.assert_array_almost_equal(
-        dx.contact.dist, np.array([-0.15, -0.15])
+        dx._impl.contact.dist, np.array([-0.15, -0.15])
     )
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[:, 0], np.array([[0, 0, -1]] * 2)
+        dx._impl.contact.frame[:, 0], np.array([[0, 0, -1]] * 2)
     )
 
   _CAP_EDGE_BOX = """
@@ -453,11 +457,11 @@ class CapsuleCollisionTest(parameterized.TestCase):
     """Tests edge contact."""
     d, dx = _collide(self._CAP_EDGE_BOX)
 
-    c = dx.contact
+    c = dx._impl.contact
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
@@ -469,13 +473,17 @@ class CapsuleCollisionTest(parameterized.TestCase):
     )
     _, dx = _collide(xml)
 
-    np.testing.assert_array_equal(dx.contact.dist < 0, np.array([True, False]))
+    np.testing.assert_array_equal(
-    np.testing.assert_array_almost_equal(dx.contact.dist[0], np.array([-0.13]))
+        dx._impl.contact.dist < 0, np.array([True, False])
-    np.testing.assert_array_almost_equal(
-        dx.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3
     )
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3
+        dx._impl.contact.dist[0], np.array([-0.13])
+    )
+    np.testing.assert_array_almost_equal(
+        dx._impl.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3
+    )
+    np.testing.assert_array_almost_equal(
+        dx._impl.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3
     )
 
   def test_capsule_convex_edge_shallow_tip(self):
@@ -489,16 +497,16 @@ class CapsuleCollisionTest(parameterized.TestCase):
     xml = xml.replace('<body pos="0.5 0 0.55"', '<body pos="0.58 0 0.55"')
     d, dx = _collide(xml)
 
-    c = dx.contact
+    c = dx._impl.contact
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip1', 1e-4)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[0][0, :3], np.array([-0.43952, 0.0, -0.898233])
+        dx._impl.contact.frame[0][0, :3], np.array([-0.43952, 0.0, -0.898233])
     )
 
     # the capsule sphere is outside the edge voronoi region, so there is a
@@ -510,16 +518,16 @@ class CapsuleCollisionTest(parameterized.TestCase):
     xml = xml.replace('<body pos="0.5 0 0.55"', '<body pos="0.5 0 0.52"')
     d, dx = _collide(xml)
 
-    c = dx.contact
+    c = dx._impl.contact
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip2', 1e-4)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[0][0, :3], np.array([0.0, 0.0, -1.0])
+        dx._impl.contact.frame[0][0, :3], np.array([0.0, 0.0, -1.0])
     )
 
 
@@ -542,17 +550,17 @@ class CylinderTest(absltest.TestCase):
     d, dx = _collide(self._CYLINDER_PLANE)
 
     # cylinder is lying flat
-    np.testing.assert_array_less(dx.contact.dist[:2], 0)
+    np.testing.assert_array_less(dx._impl.contact.dist[:2], 0)
-    np.testing.assert_array_less(-dx.contact.dist[2:], 0)
+    np.testing.assert_array_less(-dx._impl.contact.dist[2:], 0)
 
     # sort position for comparison
-    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
-    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
     idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
     d.contact.pos[:] = d.contact.pos[idx]
 
     # extract the contact points with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:2], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:2], dx._impl.contact)
     c = c.replace(dim=c.dim[:2], efc_address=c.efc_address[:2])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'cylinder_plane', 1e-5)
@@ -564,9 +572,11 @@ class CylinderTest(absltest.TestCase):
     xml = xml.replace('pos="0 0 0.04"', 'pos="0 0 0.095"')
     d, dx = _collide(xml)
 
-    np.testing.assert_array_less(dx.contact.dist, 0)
+    np.testing.assert_array_less(dx._impl.contact.dist, 0)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'cylinder_plane', 1e-5)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'cylinder_plane', 1e-5
+      )
 
   _SPHERE_CYLINDER = """
     <mujoco>
@@ -586,10 +596,10 @@ class CylinderTest(absltest.TestCase):
     """Tests sphere cylinder contact."""
     d, dx = _collide(self._SPHERE_CYLINDER)
     d.contact.pos[0][2] = 0.05  # MJX finds the deepest point on the surface
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
+          dx._impl.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
       )
 
 
@@ -612,11 +622,11 @@ class ConvexTest(absltest.TestCase):
     """Tests box collision with a plane."""
     d, dx = _collide(self._BOX_PLANE)
 
-    np.testing.assert_array_less(dx.contact.dist[:2], 0)
+    np.testing.assert_array_less(dx._impl.contact.dist[:2], 0)
-    np.testing.assert_array_less(-dx.contact.dist[2:], 0)
+    np.testing.assert_array_less(-dx._impl.contact.dist[2:], 0)
     # extract the contact points with penetration
     c = jax.tree_util.tree_map(
-        lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact
+        lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx._impl.contact
     )
     c = c.replace(dim=c.dim[[0, 1]], efc_address=c.efc_address[[0, 1]])
     for field in dataclasses.fields(Contact):
@@ -638,16 +648,18 @@ class ConvexTest(absltest.TestCase):
     """Tests box collision with a plane."""
     d, dx = _collide(self._FLAT_BOX_PLANE)
 
-    np.testing.assert_array_less(dx.contact.dist, 0)
+    np.testing.assert_array_less(dx._impl.contact.dist, 0)
 
     # sort positions for comparison
-    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
-    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
     idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
     d.contact.pos[:] = d.contact.pos[idx]
 
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'flat_box_plane', 1e-5)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'flat_box_plane', 1e-5
+      )
 
   _BOX_BOX = """
     <mujoco>
@@ -669,7 +681,7 @@ class ConvexTest(absltest.TestCase):
   def test_box_box(self):
     """Tests a face contact for a box-box collision."""
     d, dx = _collide(self._BOX_BOX, keyframe=0)
-    c = dx.contact
+    c = dx._impl.contact
 
     self.assertEqual(c.pos.shape[0], 4)
     np.testing.assert_array_less(c.dist, 0)
@@ -701,10 +713,10 @@ class ConvexTest(absltest.TestCase):
     d, dx = _collide(self._BOX_BOX_EDGE)
 
     # Only one contact point.
-    np.testing.assert_array_less(dx.contact.dist[:1], 0)
+    np.testing.assert_array_less(dx._impl.contact.dist[:1], 0)
-    np.testing.assert_array_less(-dx.contact.dist[1:], 0)
+    np.testing.assert_array_less(-dx._impl.contact.dist[1:], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
@@ -736,7 +748,7 @@ class ConvexTest(absltest.TestCase):
         ).read_bytes(),
     }
     _, dx = _collide(self._CONVEX_CONVEX, assets=assets)
-    c = dx.contact
+    c = dx._impl.contact
 
     # Only one contact point for an edge contact.
     self.assertLess(c.dist[0], 0)
@@ -766,7 +778,7 @@ class ConvexTest(absltest.TestCase):
   def test_convex_convex_edge(self):
     """Tests convex-convex collisions with edge contact via _sat_gaussmap."""
     _, dx = _collide(self._CONVEX_CONVEX_THIN)
-    c = dx.contact
+    c = dx._impl.contact
 
     # Only one contact point for an edge contact.
     self.assertLess(c.dist[0], 0)
@@ -781,7 +793,7 @@ class ConvexTest(absltest.TestCase):
             'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'
         )
     )
-    c = dx.contact
+    c = dx._impl.contact
     self.assertTrue((c.dist > 0).all())
 
 
@@ -837,8 +849,8 @@ class HFieldTest(absltest.TestCase):
 
     # check that all geoms are colliding with the hfield
     for geom_id in [1, 2, 3]:
-      mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1)
+      mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1)
-      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact)
+      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact)
       self.assertTrue((c.dist < 0).any())
       self.assertTrue((c.dist > -1e-3).any())
       # all contact normals are roughly pointing in the right direction
@@ -850,7 +862,7 @@ class HFieldTest(absltest.TestCase):
     for p in positions:
       xml = self._HFIELD.replace('<body pos="0 0', f'<body pos="{p}')
       _, dx = _collide(xml)
-      self.assertTrue((dx.contact.dist >= 0).all())
+      self.assertTrue((dx._impl.contact.dist >= 0).all())
 
   def test_hfield_deep(self):
     """Tests that objects with deep penetration do not get stuck."""
@@ -868,8 +880,8 @@ class HFieldTest(absltest.TestCase):
 
     # check that all geoms are colliding with the hfield
     for geom_id in [1, 2, 3]:
-      mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1)
+      mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1)
-      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact)
+      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact)
       # all contact normals are in the top half-face of the hfield
       self.assertTrue((c.frame[:, 0].dot(np.array([0, 0, 1])) > 0.7).all())
 
@@ -892,8 +904,8 @@ class BodyPairFilterTest(absltest.TestCase):
   def test_filter_self_collision(self):
     """Tests that self collisions get filtered."""
     d, dx = _collide(self._SELF_COLLISION)
-    self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
+    self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
-    self.assertEqual(dx.contact.pos.shape[0], 0)
+    self.assertEqual(dx._impl.contact.pos.shape[0], 0)
 
   _PARENT_CHILD = """
     <mujoco>
@@ -923,8 +935,8 @@ class BodyPairFilterTest(absltest.TestCase):
     dx = kinematics_jit_fn(mx, dx)
     dx = collision_jit_fn(mx, dx)
 
-    self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
+    self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
-    self.assertEqual(dx.contact.pos.shape[0], 0)
+    self.assertEqual(dx._impl.contact.pos.shape[0], 0)
 
   def test_disable_filter_parent_child(self):
     """Tests that filterparent flag disables parent-child filtering."""
@@ -941,8 +953,8 @@ class BodyPairFilterTest(absltest.TestCase):
     dx = collision_jit_fn(mx, dx)
 
     # one collision between parent-child spheres
-    self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
+    self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
-    self.assertEqual(dx.contact.pos.shape[0], 1)
+    self.assertEqual(dx._impl.contact.pos.shape[0], 1)
 
 
 class DimTest(parameterized.TestCase):
@@ -1027,8 +1039,8 @@ class TopKContactTest(absltest.TestCase):
     dx_all = collision_jit_fn(mx_all, dx)
     dx_top_k = collision_jit_fn(mx_top_k, dx)
 
-    self.assertEqual(dx_all.contact.dist.shape, (3,))
+    self.assertEqual(dx_all._impl.contact.dist.shape, (3,))
-    self.assertEqual(dx_top_k.contact.dist.shape, (2,))
+    self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,))
 
   _CAPSULES_MAX_PAIR = """
     <mujoco>
@@ -1073,9 +1085,9 @@ class TopKContactTest(absltest.TestCase):
     dx_all = collision_jit_fn(mx_all, dx)
     dx_top_k = collision_jit_fn(mx_top_k, dx)
 
-    self.assertEqual(dx_all.contact.dist.shape, (6,))
+    self.assertEqual(dx_all._impl.contact.dist.shape, (6,))
-    self.assertEqual(dx_top_k.contact.dist.shape, (2,))
+    self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,))
-    self.assertTrue((dx_top_k.contact.dist < 0).all())
+    self.assertTrue((dx_top_k._impl.contact.dist < 0).all())
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/collision_sdf.py

@@ -32,7 +32,9 @@ from mujoco.mjx._src.collision_types import Collision
 from mujoco.mjx._src.collision_types import GeomInfo
 from mujoco.mjx._src.dataclasses import PyTreeNode
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 # pylint: enable=g-importing-member
 
 # the SDF function takes position in, and returns a distance or objective

mjx/mujoco/mjx/_src/constraint.py

@@ -28,10 +28,12 @@ from mujoco.mjx._src.types import ConeType
 from mujoco.mjx._src.types import ConstraintType
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import EqType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import ObjType
 # pylint: enable=g-importing-member
 import numpy as np
@@ -281,6 +283,10 @@ def _efc_equality_joint(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for tendon equality constraints."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_equality_tendon requires JAX backend implementation.'
+    )
 
   eq_id = np.nonzero(m.eq_type == EqType.TENDON)[0]
 
@@ -314,9 +320,9 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
     return jax.tree_util.tree_map(lambda x: x * active, efc)
 
   inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id]
-  jac1, jac2 = d.ten_J[obj1id], d.ten_J[obj2id]
+  jac1, jac2 = d._impl.ten_J[obj1id], d._impl.ten_J[obj2id]
-  pos1 = d.ten_length[obj1id] - m.tendon_length0[obj1id]
+  pos1 = d._impl.ten_length[obj1id] - m.tendon_length0[obj1id]
-  pos2 = d.ten_length[obj2id] - m.tendon_length0[obj2id]
+  pos2 = d._impl.ten_length[obj2id] - m.tendon_length0[obj2id]
   invweight = inv1 + inv2 * (obj2id > -1)
 
   return rows(
@@ -326,8 +332,11 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_friction(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for dof frictionloss."""
-  dof_id = np.nonzero(m.dof_hasfrictionloss)[0]
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
-  tendon_id = np.nonzero(m.tendon_hasfrictionloss)[0]
+    raise ValueError('_efc_friction requires JAX backend implementation.')
+
+  dof_id = np.nonzero(m._impl.dof_hasfrictionloss)[0]
+  tendon_id = np.nonzero(m._impl.tendon_hasfrictionloss)[0]
 
   size = dof_id.size + tendon_id.size
   if (m.opt.disableflags & DisableBit.FRICTIONLOSS) or (size == 0):
@@ -337,7 +346,7 @@ def _efc_friction(m: Model, d: Data) -> Optional[_Efc]:
   args_dof += (m.dof_solimp,)
   args_dof = jax.tree_util.tree_map(lambda x: x[dof_id], args_dof)
 
-  args_ten = (d.ten_J, m.tendon_frictionloss, m.tendon_invweight0)
+  args_ten = (d._impl.ten_J, m.tendon_frictionloss, m.tendon_invweight0)
   args_ten += (m.tendon_solref_fri, m.tendon_solimp_fri)
   args_ten = jax.tree_util.tree_map(lambda x: x[tendon_id], args_ten)
 
@@ -415,6 +424,9 @@ def _efc_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for tendon limits."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_efc_limit_tendon requires JAX backend implementation.')
+
   tendon_id = np.nonzero(m.tendon_limited)[0]
 
   if (m.opt.disableflags & DisableBit.LIMIT) or tendon_id.size == 0:
@@ -423,8 +435,8 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
   length, j, range_, margin, invweight, solref, solimp = jax.tree_util.tree_map(
       lambda x: x[tendon_id],
       (
-          d.ten_length,
+          d._impl.ten_length,
-          d.ten_J,
+          d._impl.ten_J,
           m.tendon_range,
           m.tendon_margin,
           m.tendon_invweight0,
@@ -446,8 +458,12 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for frictionless contacts."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_contact_frictionless requires JAX backend implementation.'
+    )
 
-  con_id = np.nonzero(d.contact.dim == 1)[0]
+  con_id = np.nonzero(d._impl.contact.dim == 1)[0]
 
   if con_id.size == 0:
     return None
@@ -473,15 +489,19 @@ def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
         jp.zeros_like(pos),
     )
 
-  contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
+  contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
 
   return rows(contact)
 
 
 def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]:
   """Calculates constraint rows for frictional pyramidal contacts."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_contact_pyramidal requires JAX backend implementation.'
+    )
 
-  con_id = np.nonzero(d.contact.dim == condim)[0]
+  con_id = np.nonzero(d._impl.contact.dim == condim)[0]
 
   if con_id.size == 0:
     return None
@@ -518,15 +538,19 @@ def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]:
         jp.zeros_like(pos),
     )
 
-  contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
+  contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
   # concatenate to drop row grouping
   return jax.tree_util.tree_map(jp.concatenate, rows(contact))
 
 
 def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
   """Calculates constraint rows for frictional elliptic contacts."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_contact_elliptic requires JAX backend implementation.'
+    )
 
-  con_id = np.nonzero(d.contact.dim == condim)[0]
+  con_id = np.nonzero(d._impl.contact.dim == condim)[0]
 
   if con_id.size == 0:
     return None
@@ -563,7 +587,7 @@ def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
         jp.zeros_like(pos),
     )
 
-  contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
+  contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
   # concatenate to drop row grouping
   return jax.tree_util.tree_map(jp.concatenate, rows(contact))
 
@@ -602,14 +626,14 @@ def make_efc_type(
 
   if not m.opt.disableflags & DisableBit.FRICTIONLOSS:
     nf_dof = (
-        m.dof_hasfrictionloss.sum()
+        m._impl.dof_hasfrictionloss.sum()
-        if isinstance(m, Model)
+        if isinstance(m, Model) and isinstance(m._impl, ModelJAX)
         else (m.dof_frictionloss > 0).sum()
     )
     efc_types += [ConstraintType.FRICTION_DOF] * nf_dof
     nf_tendon = (
-        m.tendon_hasfrictionloss.sum()
+        m._impl.tendon_hasfrictionloss.sum()
-        if isinstance(m, Model)
+        if isinstance(m, Model) and isinstance(m._impl, ModelJAX)
         else (m.tendon_frictionloss > 0).sum()
     )
     efc_types += [ConstraintType.FRICTION_TENDON] * nf_tendon
@@ -683,10 +707,14 @@ def make_constraint(m: Model, d: Data) -> Data:
 
   if not efcs:
     z = jp.empty(0)
-    d = d.replace(efc_J=jp.empty((0, m.nv)))
+    d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))})
-    d = d.replace(
+    d = d.tree_replace({
-        efc_D=z, efc_aref=z, efc_frictionloss=z, efc_pos=z, efc_margin=z
+        '_impl.efc_D': z,
-    )
+        '_impl.efc_aref': z,
+        '_impl.efc_frictionloss': z,
+        '_impl.efc_pos': z,
+        '_impl.efc_margin': z,
+    })
     return d
 
   efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
@@ -699,9 +727,13 @@ def make_constraint(m: Model, d: Data) -> Data:
     return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss
 
   aref, r, pos, margin, frictionloss = fn(efc)
-  d = d.replace(
+  d = d.tree_replace({
-      efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos, efc_margin=margin
+      '_impl.efc_J': efc.J,
-  )
+      '_impl.efc_D': 1 / r,
-  d = d.replace(efc_frictionloss=frictionloss)
+      '_impl.efc_aref': aref,
+      '_impl.efc_pos': pos,
+      '_impl.efc_margin': margin,
+  })
+  d = d.tree_replace({'_impl.efc_frictionloss': frictionloss})
 
   return d

mjx/mujoco/mjx/_src/constraint_test.py

@@ -69,16 +69,16 @@ class ConstraintTest(parameterized.TestCase):
 
       order = test_util.efc_order(m, d, dx)
       d_efc_j = d.efc_J.reshape((-1, m.nv))
-      _assert_eq(d_efc_j, dx.efc_J[order][: d.nefc], 'efc_J')
+      _assert_eq(d_efc_j, dx._impl.efc_J[order][: d.nefc], 'efc_J')
-      _assert_eq(0, dx.efc_J[order][d.nefc :], 'efc_J')
+      _assert_eq(0, dx._impl.efc_J[order][d.nefc :], 'efc_J')
-      _assert_eq(d.efc_aref, dx.efc_aref[order][: d.nefc], 'efc_aref')
+      _assert_eq(d.efc_aref, dx._impl.efc_aref[order][: d.nefc], 'efc_aref')
-      _assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref')
+      _assert_eq(0, dx._impl.efc_aref[order][d.nefc :], 'efc_aref')
-      _assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D')
+      _assert_eq(d.efc_D, dx._impl.efc_D[order][: d.nefc], 'efc_D')
-      _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
+      _assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos')
-      _assert_eq(dx.efc_pos[order][d.nefc :], 0, 'efc_pos')
+      _assert_eq(dx._impl.efc_pos[order][d.nefc :], 0, 'efc_pos')
       _assert_eq(
           d.efc_frictionloss,
-          dx.efc_frictionloss[order][: d.nefc],
+          dx._impl.efc_frictionloss[order][: d.nefc],
           'efc_frictionloss',
       )
 
@@ -104,7 +104,7 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 0)
     self.assertEqual(nc, 0)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 0)
+    self.assertEqual(dx._impl.efc_J.shape[0], 0)
 
   def test_disable_equality(self):
     m = test_util.load_test_file('constraints.xml')
@@ -115,7 +115,9 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 180)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 187)  # only joint/tendon limit, contact
+    self.assertEqual(
+        dx._impl.efc_J.shape[0], 187
+    )  # only joint/tendon limit, contact
 
   def test_disable_contact(self):
     m = test_util.load_test_file('constraints.xml')
@@ -126,7 +128,9 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 0)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 27)  # only equality, joint/tendon limit
+    self.assertEqual(
+        dx._impl.efc_J.shape[0], 27
+    )  # only equality, joint/tendon limit
 
   def test_disable_frictionloss(self):
     m = test_util.load_test_file('constraints.xml')
@@ -137,7 +141,7 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 180)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 205)
+    self.assertEqual(dx._impl.efc_J.shape[0], 205)
 
   def test_margin(self):
     """Test margin."""
@@ -163,8 +167,8 @@ class ConstraintTest(parameterized.TestCase):
     dx = mjx.make_constraint(mx, dx)
 
     order = test_util.efc_order(m, d, dx)
-    _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
+    _assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos')
-    _assert_eq(d.efc_margin, dx.efc_margin[order][: d.nefc], 'efc_margin')
+    _assert_eq(d.efc_margin, dx._impl.efc_margin[order][: d.nefc], 'efc_margin')
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/dataclasses.py

@@ -18,6 +18,7 @@ import copy
 import dataclasses
 import typing
 from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
+import warnings
 import jax
 import numpy as np
 
@@ -49,6 +50,36 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
   """
   data_clz = dataclasses.dataclass(frozen=True)(clz)
   data_clz.replace = dataclasses.replace
+  # def replace(self, **updates):
+  #   """Returns a new object replacing the specified fields with new values."""
+  #   if not hasattr(self, '_impl'):
+  #     return dataclasses.replace(self, **updates)
+
+  #   # Private fields under `_impl` are allowed to be replaced directly as if
+  #   # they were on the base class. This logic will be removed in a future
+  #   # release.
+  #   impl_updates = {}
+  #   for k in tuple(updates.keys()):
+  #     # Recall that getattr is overridden for '_impl' fields.
+  #     hasattr_ = k in self.__annotations__
+  #     if not hasattr_ and hasattr(self._impl, k):  # pylint: disable=protected-access
+  #       impl_updates[k] = updates[k]
+  #       del updates[k]
+
+  #   if impl_updates:
+  #     updates['_impl'] = self._impl.replace(**impl_updates)
+  #     warnings.warn(
+  #         f'Accessing/replacing fields `{tuple(impl_updates.keys())}` directly'
+  #         f' from `{self.__class__.__name__}` will be deprecated. Refrain from'
+  #         ' using private fields that were moved to'
+  #         f' `{self.__class__.__name__}`._impl.',
+  #         DeprecationWarning,
+  #         stacklevel=2,
+  #     )
+
+  #   return dataclasses.replace(self, **updates)
+
+  # data_clz.replace = replace
 
   if register_as_pytree:
     meta_fields, data_fields = [], []

mjx/mujoco/mjx/_src/derivative.py

@@ -40,7 +40,9 @@ def deriv_smooth_vel(m: Model, d: Data) -> Optional[jax.Array]:
     gain_vel = m.actuator_gainprm[:, 2] * affine_gain
     ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act)
     vel = bias_vel + gain_vel * ctrl
-    qderiv = d.actuator_moment.T @ jax.vmap(jp.multiply)(d.actuator_moment, vel)
+    qderiv = d._impl.actuator_moment.T @ jax.vmap(jp.multiply)(
+        d._impl.actuator_moment, vel
+    )
 
   # qDeriv += d qfrc_passive / d qvel
   if not m.opt.disableflags & DisableBit.PASSIVE:
@@ -49,9 +51,9 @@ def deriv_smooth_vel(m: Model, d: Data) -> Optional[jax.Array]:
     else:
       qderiv -= jp.diag(m.dof_damping)
     if m.ntendon:
-      qderiv -= d.ten_J.T @ jp.diag(m.tendon_damping) @ d.ten_J
+      qderiv -= d._impl.ten_J.T @ jp.diag(m.tendon_damping) @ d._impl.ten_J
     # TODO(robotics-simulation): fluid drag model
-    if m.opt.has_fluid_params:
+    if m.opt.has_fluid_params:  # pytype: disable=attribute-error
       raise NotImplementedError('fluid drag not supported for implicitfast')
 
   # TODO(team): rne derivative

mjx/mujoco/mjx/_src/forward.py

@@ -33,12 +33,14 @@ from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import BiasType
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import DynType
 from mujoco.mjx._src.types import GainType
 from mujoco.mjx._src.types import IntegratorType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import TrnType
 # pylint: enable=g-importing-member
 import numpy as np
@@ -81,10 +83,13 @@ def fwd_position(m: Model, d: Data) -> Data:
 @named_scope
 def fwd_velocity(m: Model, d: Data) -> Data:
   """Velocity-dependent computations."""
-  d = d.replace(
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
-      actuator_velocity=d.actuator_moment @ d.qvel,
+    raise ValueError('fwd_velocity requires JAX backend implementation.')
-      ten_velocity=d.ten_J @ d.qvel,
+
-  )
+  d = d.tree_replace({
+      '_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel,
+      '_impl.ten_velocity': d._impl.ten_J @ d.qvel,
+  })
   d = smooth.com_vel(m, d)
   d = passive.passive(m, d)
   d = smooth.rne(m, d)
@@ -173,8 +178,8 @@ def fwd_actuation(m: Model, d: Data) -> Data:
       m.actuator_gainprm,
       m.actuator_biastype,
       m.actuator_biasprm,
-      d.actuator_length,
+      d._impl.actuator_length,
-      d.actuator_velocity,
+      d._impl.actuator_velocity,
       ctrl_act,
       jp.array(m.actuator_lengthrange),
       jp.array(m.actuator_acc0),
@@ -215,7 +220,7 @@ def fwd_actuation(m: Model, d: Data) -> Data:
   )
   force = jp.clip(force, forcerange[:, 0], forcerange[:, 1])
 
-  qfrc_actuator = d.actuator_moment.T @ force
+  qfrc_actuator = d._impl.actuator_moment.T @ force
 
   if m.ngravcomp:
     # actuator-level gravity compensation, skip if added as passive force
@@ -331,13 +336,17 @@ def _advance(
 @named_scope
 def euler(m: Model, d: Data) -> Data:
   """Euler integrator, semi-implicit in velocity."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('euler requires JAX backend implementation.')
+
   # integrate damping implicitly
   qacc = d.qacc
   if not m.opt.disableflags & DisableBit.EULERDAMP:
     if support.is_sparse(m):
-      dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
+      qM = d._impl.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
     else:
-      dh = d.replace(qM=d.qM + jp.diag(m.opt.timestep * m.dof_damping))
+      qM = d._impl.qM + jp.diag(m.opt.timestep * m.dof_damping)
+    dh = d.tree_replace({'_impl.qM': qM})
     dh = smooth.factor_m(m, dh)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
     qacc = smooth.solve_m(m, dh, qfrc)
@@ -392,13 +401,15 @@ def rungekutta4(m: Model, d: Data) -> Data:
 @named_scope
 def implicit(m: Model, d: Data) -> Data:
   """Integrates fully implicit in velocity."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('implicit requires JAX backend implementation.')
 
   qderiv = derivative.deriv_smooth_vel(m, d)
 
   qacc = d.qacc
   if qderiv is not None:
     # TODO(robotics-simulation): use smooth.factor_m / solve_m here:
-    qm = support.full_m(m, d) if support.is_sparse(m) else d.qM
+    qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.qM
     qm -= m.opt.timestep * qderiv
     qh, _ = jax.scipy.linalg.cho_factor(qm)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
@@ -410,6 +421,9 @@ def implicit(m: Model, d: Data) -> Data:
 @named_scope
 def forward(m: Model, d: Data) -> Data:
   """Forward dynamics."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('forward requires JAX backend implementation.')
+
   d = fwd_position(m, d)
   d = sensor.sensor_pos(m, d)
   d = fwd_velocity(m, d)
@@ -418,7 +432,7 @@ def forward(m: Model, d: Data) -> Data:
   d = fwd_acceleration(m, d)
   d = sensor.sensor_acc(m, d)
 
-  if d.efc_J.size == 0:
+  if d._impl.efc_J.size == 0:
     d = d.replace(qacc=d.qacc_smooth)
     return d
 

mjx/mujoco/mjx/_src/inverse.py

@@ -69,16 +69,16 @@ def inv_constraint(m: Model, d: Data) -> Data:
   """Inverse constraint solver."""
 
   # no constraints
-  if d.efc_J.size == 0:
+  if d._impl.efc_J.size == 0:  # pytype: disable=attribute-error
     return d.replace(qfrc_constraint=jp.zeros(m.nv))
 
   # update
   ctx = solver.Context.create(m, d, grad=False)
 
-  return d.replace(
+  return d.tree_replace({
-      qfrc_constraint=ctx.qfrc_constraint,
+      'qfrc_constraint': ctx.qfrc_constraint,
-      efc_force=ctx.efc_force,
+      '_impl.efc_force': ctx.efc_force,
-  )
+  })
 
 
 def inverse(m: Model, d: Data) -> Data:

mjx/mujoco/mjx/_src/inverse_test.py

@@ -15,12 +15,10 @@
 """Tests for inverse dynamics functions."""
 from absl.testing import absltest
 from absl.testing import parameterized
-import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import support
-from mujoco.mjx._src import test_util
 import numpy as np
 
 # tolerance for difference between MuJoCo and MJX calculations - mostly

mjx/mujoco/mjx/_src/io_test.py

@@ -14,14 +14,19 @@
 # ==============================================================================
 """Tests for io functions."""
 
+import os
+from unittest import mock
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import io as mjx_io
 from mujoco.mjx._src import test_util
+
 # pylint: disable=g-importing-member
+from mujoco.mjx._src.types import BackendImpl
 from mujoco.mjx._src.types import ConeType
 # pylint: enable=g-importing-member
 import numpy as np
@@ -95,10 +100,13 @@ _MULTIPLE_CONSTRAINTS = """
 class ModelIOTest(parameterized.TestCase):
   """IO tests for mjx.Model."""
 
-  @parameterized.parameters(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS)
+  @parameterized.product(
-  def test_put_model(self, xml):
+      xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS),
+      backend_impl=('jax', 'c'),
+  )
+  def test_put_model(self, xml, backend_impl):
     m = mujoco.MjModel.from_xml_string(xml)
-    mx = mjx.put_model(m)
+    mx = mjx.put_model(m, backend_impl=backend_impl)
 
     def assert_not_weak_type(x):
       if isinstance(x, jax.Array):
@@ -120,8 +128,14 @@ class ModelIOTest(parameterized.TestCase):
     self.assertEqual(mx.nM, m.nM)
     self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
 
-    # fields restricted to MuJoCo should not be populated
+    if backend_impl == 'jax':
-    self.assertEqual(mx.bvh_aabb.shape, (0,))
+      # fields restricted to MuJoCo should not be populated
+      self.assertFalse(hasattr(mx, 'bvh_aabb'))
+    elif backend_impl == 'c':
+      # Options specific to C are populated.
+      self.assertEqual(mx.opt.apirate, m.opt.apirate)
+      # Fields private to C backend impl are populated.
+      self.assertTrue(hasattr(mx._impl, 'bvh_aabb'))
 
     np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
     np.testing.assert_allclose(mx.geom_type, m.geom_type)
@@ -150,7 +164,8 @@ class ModelIOTest(parameterized.TestCase):
     m = mjx.put_model(
         mujoco.MjModel.from_xml_string(
             '<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
-        )
+        ),
+        backend_impl='jax',
     )
     self.assertTrue(m.opt.has_fluid_params)
 
@@ -174,7 +189,8 @@ class ModelIOTest(parameterized.TestCase):
 
   def test_margin_gap_mesh_not_implemented(self):
     with self.assertRaises(NotImplementedError):
-      mjx.put_model(mujoco.MjModel.from_xml_string("""
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string("""
         <mujoco>
           <asset>
             <mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 1"/>
@@ -189,27 +205,32 @@ class ModelIOTest(parameterized.TestCase):
               <geom size="0.05"/>
             </body>
           </worldbody>
-        </mujoco>"""))
+        </mujoco>"""),
+          backend_impl='jax',
+      )
 
   def test_implicitfast_fluid_not_implemented(self):
     with self.assertRaises(NotImplementedError):
-      mjx.put_model(mujoco.MjModel.from_xml_string("""
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string("""
         <mujoco>
           <option viscosity="3.0" integrator="implicitfast"/>
           <worldbody/>
-        </mujoco>"""))
+        </mujoco>"""),
+          backend_impl='jax',
+      )
 
   def test_wrap_inside(self):
     m = test_util.load_test_file('tendon/wrap_sidesite.xml')
-    mx0 = mjx.put_model(m)
+    mx0 = mjx.put_model(m, backend_impl='jax')
     np.testing.assert_equal(
-        mx0.is_wrap_inside,
+        mx0._impl.is_wrap_inside,
         np.array([1, 0, 1, 0, 1, 1, 0]),
     )
     m.site_pos[2] = m.site_pos[1]
-    mx1 = mjx.put_model(m)
+    mx1 = mjx.put_model(m, backend_impl='jax')
     np.testing.assert_equal(
-        mx1.is_wrap_inside,
+        mx1._impl.is_wrap_inside,
         np.array([0, 0, 1, 0, 1, 0, 0]),
     )
 
@@ -217,11 +238,11 @@ class ModelIOTest(parameterized.TestCase):
 class DataIOTest(parameterized.TestCase):
   """IO tests for mjx.Data."""
 
-  def test_make_data(self):
+  @parameterized.parameters('jax', 'c')
+  def test_make_data(self, backend_impl: str):
     """Test that make_data returns the correct shapes."""
-
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
-    d = mjx.make_data(m)
+    d = mjx.make_data(m, backend_impl=backend_impl)
 
     nq = 22
     nbody = 5
@@ -230,7 +251,7 @@ class DataIOTest(parameterized.TestCase):
     nv = 19
     nefc = 185
 
-    self.assertEqual(d.nefc, nefc)
+    self.assertEqual(d._impl.nefc, nefc)
     self.assertEqual(d.qpos.shape, (nq,))
     self.assertEqual(d.qvel.shape, (nv,))
     self.assertEqual(d.act.shape, (0,))
@@ -251,57 +272,68 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d.geom_xpos.shape, (6, 3))
     self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
     self.assertEqual(d.subtree_com.shape, (nbody, 3))
-    self.assertEqual(d.cdof.shape, (nv, 6))
+    self.assertEqual(d._impl.cdof.shape, (nv, 6))
-    self.assertEqual(d.cinert.shape, (nbody, 10))
+    self.assertEqual(d._impl.cinert.shape, (nbody, 10))
-    self.assertEqual(d.crb.shape, (nbody, 10))
+    self.assertEqual(d._impl.crb.shape, (nbody, 10))
-    self.assertEqual(d.actuator_length.shape, (1,))
+    self.assertEqual(d._impl.actuator_length.shape, (1,))
-    self.assertEqual(d.actuator_moment.shape, (1, nv))
+    self.assertEqual(d._impl.actuator_moment.shape, (1, nv))
-    self.assertEqual(d.qM.shape, (nv, nv))
+    self.assertEqual(d._impl.contact.dist.shape, (ncon,))
-    self.assertEqual(d.qLD.shape, (nv, nv))
+    self.assertEqual(d._impl.contact.pos.shape, (ncon, 3))
-    self.assertEqual(d.qLDiagInv.shape, (0,))
+    self.assertEqual(d._impl.contact.frame.shape, (ncon, 3, 3))
-    self.assertEqual(d.contact.dist.shape, (ncon,))
+    self.assertEqual(d._impl.contact.solref.shape, (ncon, 2))
-    self.assertEqual(d.contact.pos.shape, (ncon, 3))
+    self.assertEqual(d._impl.contact.solimp.shape, (ncon, 5))
-    self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
+    self.assertEqual(d._impl.contact.geom1.shape, (ncon,))
-    self.assertEqual(d.contact.solref.shape, (ncon, 2))
+    self.assertEqual(d._impl.contact.geom2.shape, (ncon,))
-    self.assertEqual(d.contact.solimp.shape, (ncon, 5))
+    self.assertEqual(d._impl.efc_J.shape, (nefc, nv))
-    self.assertEqual(d.contact.geom1.shape, (ncon,))
+    self.assertEqual(d._impl.efc_frictionloss.shape, (nefc,))
-    self.assertEqual(d.contact.geom2.shape, (ncon,))
+    self.assertEqual(d._impl.efc_D.shape, (nefc,))
-    self.assertEqual(d.efc_J.shape, (nefc, nv))
+    self.assertEqual(d._impl.actuator_velocity.shape, (1,))
-    self.assertEqual(d.efc_frictionloss.shape, (nefc,))
-    self.assertEqual(d.efc_D.shape, (nefc,))
-    self.assertEqual(d.actuator_velocity.shape, (1,))
     self.assertEqual(d.cvel.shape, (nbody, 6))
-    self.assertEqual(d.cdof_dot.shape, (nv, 6))
+    self.assertEqual(d._impl.cdof_dot.shape, (nv, 6))
     self.assertEqual(d.qfrc_bias.shape, (nv,))
     self.assertEqual(d.qfrc_passive.shape, (nv,))
-    self.assertEqual(d.efc_aref.shape, (nefc,))
+    self.assertEqual(d._impl.efc_aref.shape, (nefc,))
     self.assertEqual(d.qfrc_actuator.shape, (nv,))
     self.assertEqual(d.qfrc_smooth.shape, (nv,))
     self.assertEqual(d.qacc_smooth.shape, (nv,))
     self.assertEqual(d.qfrc_constraint.shape, (nv,))
     self.assertEqual(d.qfrc_inverse.shape, (nv,))
-    self.assertEqual(d.efc_force.shape, (nefc,))
+    self.assertEqual(d._impl.efc_force.shape, (nefc,))
+
+    if backend_impl == 'jax':
+      self.assertEqual(d._impl.qM.shape, (nv, nv))
+      self.assertEqual(d._impl.qLD.shape, (nv, nv))
+      self.assertEqual(d._impl.qLDiagInv.shape, (0,))
+    elif backend_impl == 'c':
+      self.assertEqual(d._impl.qM.shape, (nm,))
+      self.assertEqual(d._impl.qLD.shape, (nm,))
+      self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
 
     # test sparse
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
-    d = mjx.make_data(m)
+    d = mjx.make_data(m, backend_impl=backend_impl)
-    self.assertEqual(d.qM.shape, (nm,))
+    self.assertEqual(d._impl.qM.shape, (nm,))
-    self.assertEqual(d.qLD.shape, (nm,))
+    self.assertEqual(d._impl.qLD.shape, (nm,))
-    self.assertEqual(d.qLDiagInv.shape, (nv,))
+    self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
 
-  def test_put_data(self):
+    if backend_impl == 'c':
+      # check C specific fields
+      self.assertEqual(d._impl.light_xpos.shape, (m.nlight, 3))
+      self.assertEqual(d._impl.bvh_active.shape, (m.nbvh,))
+
+  @parameterized.parameters('jax', 'c')
+  def test_put_data(self, backend_impl: str):
     """Test that put_data puts the correct data for dense and sparse."""
-
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
 
     # check a few fields
     np.testing.assert_allclose(dx.qpos, d.qpos)
     np.testing.assert_allclose(dx.xpos, d.xpos)
     np.testing.assert_allclose(dx.cvel, d.cvel)
-    np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(dx._impl.cdof_dot, d.cdof_dot)
 
     # check that there are no weak types
     self.assertFalse(
@@ -312,21 +344,26 @@ class DataIOTest(parameterized.TestCase):
         )
     )
 
-    # check that qM is transformed properly
+    if backend_impl == 'jax':
-    qm = np.zeros((m.nv, m.nv), dtype=np.float64)
+      # check that qM is transformed properly
-    mujoco.mj_fullM(m, qm, d.qM)
+      qm = np.zeros((m.nv, m.nv), dtype=np.float64)
-    np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
+      mujoco.mj_fullM(m, qm, d.qM)
+      np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
+    elif backend_impl == 'c':
+      np.testing.assert_allclose(dx._impl.qM, d.qM)
+      np.testing.assert_allclose(dx._impl.qLD, d.qLD)
+      np.testing.assert_allclose(dx._impl.qLDiagInv, d.qLDiagInv)
 
     # 4 contacts, 2 for each capsule against the plane
-    self.assertEqual(dx.contact.dist.shape, (4,))
+    self.assertEqual(dx._impl.contact.dist.shape, (4,))
     self.assertEqual(d.ncon, 1)  # however only 1 contact in this step
-    np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0])
+    np.testing.assert_allclose(dx._impl.contact.dist[0], d.contact.dist[0])
-    self.assertTrue((dx.contact.dist[1:] > 0).all())
+    self.assertTrue((dx._impl.contact.dist[1:] > 0).all())
-    self.assertEqual(dx.contact.frame.shape, (4, 3, 3))
+    self.assertEqual(dx._impl.contact.frame.shape, (4, 3, 3))
     np.testing.assert_allclose(
-        dx.contact.frame[0].reshape(9), d.contact.frame[0]
+        dx._impl.contact.frame[0].reshape(9), d.contact.frame[0]
     )
-    np.testing.assert_allclose(dx.contact.frame[1:], 0)
+    np.testing.assert_allclose(dx._impl.contact.frame[1:], 0)
 
     # xmat, ximat, geom_xmat are all shape transformed
     self.assertEqual(dx.xmat.shape, (3, 3, 3))
@@ -339,57 +376,65 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
 
     # tendon data is correct
-    np.testing.assert_allclose(dx.ten_length, d.ten_length)
+    np.testing.assert_allclose(dx._impl.ten_length, d.ten_length)
-    np.testing.assert_equal(dx.ten_wrapadr, np.zeros((1,)))
+    np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,)))
-    np.testing.assert_equal(dx.ten_wrapnum, np.zeros((1,)))
+    np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,)))
-    np.testing.assert_equal(dx.wrap_obj, np.zeros((2, 2)))
+    np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2)))
-    np.testing.assert_equal(dx.wrap_xpos, np.zeros((2, 6)))
+    np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6)))
 
     # efc_ are also shape transformed and padded
-    self.assertEqual(dx.efc_J.shape, (45, 8))  # nefc, nv
+    self.assertEqual(dx._impl.efc_J.shape, (45, 8))  # nefc, nv
     d_efc_j = d.efc_J.reshape((-1, 8))
-    np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3])  # connect eq
+    np.testing.assert_allclose(dx._impl.efc_J[:3], d_efc_j[:3])  # connect eq
-    np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3])  # one active limit
+    np.testing.assert_allclose(
-    np.testing.assert_allclose(dx.efc_J[4], 0)  # one inactive limit
+        dx._impl.efc_J[3], d_efc_j[3]
-    np.testing.assert_allclose(dx.efc_J[5:15], d_efc_j[4:14])  # contact
+    )  # one active limit
-    np.testing.assert_allclose(dx.efc_J[15:], 0)  # no contact
+    np.testing.assert_allclose(dx._impl.efc_J[4], 0)  # one inactive limit
+    np.testing.assert_allclose(dx._impl.efc_J[5:15], d_efc_j[4:14])  # contact
+    np.testing.assert_allclose(dx._impl.efc_J[15:], 0)  # no contact
 
     # check another efc_ too
-    self.assertEqual(dx.efc_aref.shape, (45,))  # nefc
+    self.assertEqual(dx._impl.efc_aref.shape, (45,))  # nefc
-    np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
+    np.testing.assert_allclose(dx._impl.efc_aref[:3], d.efc_aref[:3])
-    np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
+    np.testing.assert_allclose(dx._impl.efc_aref[3], d.efc_aref[3])
-    np.testing.assert_allclose(dx.efc_aref[4], 0)
+    np.testing.assert_allclose(dx._impl.efc_aref[4], 0)
-    np.testing.assert_allclose(dx.efc_aref[5:15], d.efc_aref[4:14])
+    np.testing.assert_allclose(dx._impl.efc_aref[5:15], d.efc_aref[4:14])
-    np.testing.assert_allclose(dx.efc_aref[15:], 0)
+    np.testing.assert_allclose(dx._impl.efc_aref[15:], 0)
 
     # check sparse transform is correct
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx_sparse = mjx.put_data(m, d)
+    dx_sparse = mjx.put_data(m, d, backend_impl=backend_impl)
-    np.testing.assert_allclose(dx_sparse.efc_J, dx.efc_J, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8)
 
     # check sparse mass matrices are correct
-    np.testing.assert_allclose(dx_sparse.qM, d.qM, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8)
-    np.testing.assert_allclose(dx_sparse.qLD, d.qLD, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse._impl.qLD, d.qLD, atol=1e-8)
-    np.testing.assert_allclose(dx_sparse.qLDiagInv, d.qLDiagInv, atol=1e-8)
+    np.testing.assert_allclose(
+        dx_sparse._impl.qLDiagInv, d.qLDiagInv, atol=1e-8
+    )
 
     # check dense mass matrices are correct
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx_from_dense = mjx.put_data(m, d)
+    dx_from_dense = mjx.put_data(m, d, backend_impl=backend_impl)
-    qm = np.zeros((m.nv, m.nv))
+    if backend_impl == 'jax':
-    mujoco.mj_fullM(m, qm, d.qM)
+      qm = np.zeros((m.nv, m.nv))
-    np.testing.assert_allclose(dx_from_dense.qM, qm, atol=1e-8)
+      mujoco.mj_fullM(m, qm, d.qM)
+      np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
+    elif backend_impl == 'c':
+      np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8)
 
-  def test_get_data(self):
+  @parameterized.parameters('jax', 'c')
+  def test_get_data(self, backend_impl: str):
     """Test that get_data makes correct MjData."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     d_2: mujoco.MjData = mjx.get_data(m, dx)
 
     # check a few fields
@@ -424,6 +469,10 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
     np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
 
+    if backend_impl == 'c':
+      # check fields specific to the C implementation
+      np.testing.assert_allclose(d_2.M_rownnz, d.M_rownnz)
+
   def test_get_data_runs(self):
     xml = """
       <mujoco>
@@ -444,13 +493,14 @@ class DataIOTest(parameterized.TestCase):
     dx = mjx.put_data(m, d)
     mjx.get_data(m, dx)
 
-  def test_get_data_batched(self):
+  @parameterized.parameters('jax', 'c')
+  def test_get_data_batched(self, backend_impl):
     """Test that get_data makes correct List[MjData] for batched Data."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     # second data in batch has contact dist > 0, disables contact
     dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
     ds = mjx.get_data(m, dx_b)
@@ -460,13 +510,14 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(ds[0].ncon, 1)
     self.assertEqual(ds[1].ncon, 0)
 
-  def test_get_data_into(self):
+  @parameterized.parameters('jax', 'c')
+  def test_get_data_into(self, backend_impl):
     """Test that get_data_into correctly populates an MjData."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     d_2 = mujoco.MjData(m)
     mjx.get_data_into(d_2, m, dx)
 
@@ -482,32 +533,33 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d_2.contact.frame.shape, (1, 9))
     np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
 
-  def test_get_data_into_wrong_shape(self):
+  @parameterized.parameters('jax', 'c')
+  def test_get_data_into_wrong_shape(self, backend_impl):
     """Tests that get_data_into throwsif input and output shapes don't match."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     m_2 = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
     d_2 = mujoco.MjData(m_2)
     with self.assertRaisesRegex(ValueError, r'Input field.*has shape.*'):
       mjx.get_data_into(d_2, m, dx)
 
-  def test_make_matches_put(self):
+  @parameterized.parameters('jax', 'c')
+  def test_make_matches_put(self, backend_impl):
     """Test that make_data produces a pytree that matches put_data."""
-
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
 
     step_fn = lambda d: d.replace(time=d.time + 1)
     step_fn_jit = jax.jit(step_fn).lower(dx).compile()
 
     # placing an MjData onto device should yield the same treedef mjx.Data as
     # calling make_data.  they should be interchangeable for jax functions:
-    step_fn_jit(mjx.make_data(m))
+    step_fn_jit(mjx.make_data(m, backend_impl=backend_impl))
 
   def test_contact_elliptic_condim1(self):
     """Test that condim=1 with ConeType.ELLIPTIC is not implemented."""
@@ -556,7 +608,327 @@ class DataIOTest(parameterized.TestCase):
         </mujoco>
       """)
     with self.assertRaises(NotImplementedError):
-      mjx.put_model(m)
+      mjx.put_model(m, backend_impl='jax')
+
+
+class FullCompatTest(parameterized.TestCase):
+  """Tests for the _full_compat flag."""
+
+  def test_full_compat_deprecated(self):
+    """Tests that _full_compat is deprecated."""
+    xml = """
+      <mujoco>
+        <worldbody>
+          <body name="box">
+            <joint name="slide1" type="slide" axis="1 0 0" />
+            <geom type="box" size=".05 .05 .05" mass="1"/>
+          </body>
+        </worldbody>
+        <actuator>
+          <motor joint="slide1"/>
+        </actuator>
+      </mujoco>
+    """
+    m = mujoco.MjModel.from_xml_string(xml)
+    with self.assertWarns(DeprecationWarning):
+      out = mjx_io.put_model(m, _full_compat=True)
+      self.assertEqual(out.backend_impl, BackendImpl.C)
+    with self.assertWarns(DeprecationWarning):
+      out = mjx_io.make_data(m, _full_compat=True)
+      self.assertEqual(out.backend_impl, BackendImpl.C)
+
+
+# Test cases for `_resolve_backend_impl_and_device` where the device is
+# specified by the user and the device is available.
+_DEVICE_TEST_CASES = [
+    # Arguments use the following format:
+    # (device_type_str, backend_impl_str,
+    #  (expected_device, expected_backend_impl)))
+    # No backend specified.
+    ('cpu', None, ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
+    ('tpu', None, ('tpu', BackendImpl.JAX)),
+    # JAX backend specified.
+    ('cpu', 'jax', ('cpu', BackendImpl.JAX)),
+    ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('tpu', 'jax', ('tpu', BackendImpl.JAX)),
+    # WARP backend specified.
+    ('cpu', 'warp', ('cpu', 'error')),
+    ('gpu-notnvidia', 'warp', ('cpu', 'error')),
+    ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
+    ('tpu', 'warp', ('tpu', 'error')),
+    # C backend specified.
+    ('cpu', 'c', ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', 'c', ('cpu', 'error')),
+    ('gpu-nvidia', 'c', ('cpu', 'error')),
+    ('tpu', 'c', ('tpu', 'error')),
+]
+
+# Test cases for `_resolve_backend_impl_and_device` where the user does NOT
+# specify a device. We mock the JAX default device.
+_DEFAULT_DEVICE_TEST_CASES = [
+    # Arguments use the following format:
+    # (jax.default_device, backend_impl_str,
+    #  (expected_device, expected_backend_impl))
+    # No backend impl specified.
+    ('cpu', None, ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
+    ('tpu', None, ('tpu', BackendImpl.JAX)),
+    # JAX backend impl specified.
+    ('cpu', 'jax', ('cpu', BackendImpl.JAX)),
+    ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('tpu', 'jax', ('tpu', BackendImpl.JAX)),
+    # WARP backend impl specified.
+    ('cpu', 'warp', ('cpu', 'error')),
+    ('gpu-notnvidia', 'warp', ('cpu', 'error')),
+    ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
+    ('tpu', 'warp', ('tpu', 'error')),
+    # C backend impl specified, CPU should always be available.
+    ('cpu', 'c', ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', 'c', ('cpu', BackendImpl.C)),
+    ('gpu-nvidia', 'c', ('cpu', BackendImpl.C)),
+    ('tpu', 'c', ('cpu', BackendImpl.C)),
+]
+
+
+class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
+  """Tests for the _resolve_backend_impl_and_device function."""
+
+  def setUp(self):
+    super().setUp()
+
+    # Create mock devices
+    self.mock_cpu = mock.Mock(spec=jax.Device)
+    self.mock_cpu.platform = 'cpu'
+    self.mock_cpu.device_kind = 'Mock CPU'
+    self.mock_cpu.id = 0
+
+    self.mock_nvidia_gpu = mock.Mock(spec=jax.Device)
+    self.mock_nvidia_gpu.platform = 'gpu'
+    self.mock_nvidia_gpu.device_kind = 'NVIDIA Mocked GPU'
+    self.mock_nvidia_gpu.id = 0
+
+    self.mock_other_gpu = mock.Mock(spec=jax.Device)
+    self.mock_other_gpu.platform = 'gpu'
+    self.mock_other_gpu.device_kind = 'Other Mocked GPU'
+    self.mock_other_gpu.id = 1
+
+    self.mock_tpu = mock.Mock(spec=jax.Device)
+    self.mock_tpu.platform = 'tpu'
+    self.mock_tpu.device_kind = 'Mock TPU'
+    self.mock_tpu.id = 0
+
+    # Patch jax.devices for the entire test class using enter_context
+    self.mock_jax_devices = self.enter_context(mock.patch('jax.devices'))
+    self.mock_default_backend = self.enter_context(
+        mock.patch('jax.default_backend')
+    )
+
+  @parameterized.named_parameters(
+      (f'{str(args[0])}_{str(args[1])}', *args) for args in _DEVICE_TEST_CASES
+  )
+  @mock.patch.dict(
+      os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'}
+  )
+  def test_resolve_with_device(
+      self,
+      device_type_str,
+      backend_impl_str,
+      expected,
+  ):
+    """Tests various combinations of device and backend impls."""
+    input_device = {
+        'cpu': self.mock_cpu,
+        'gpu-nvidia': self.mock_nvidia_gpu,
+        'gpu-notnvidia': self.mock_other_gpu,
+        'tpu': self.mock_tpu,
+    }[device_type_str]
+
+    def devices_side_effect(backend=None):
+      # assume the user-specified device is always available
+      if backend == 'cpu':
+        return [self.mock_cpu]
+      elif backend == 'gpu':
+        if 'nvidia' in device_type_str:
+          return [self.mock_nvidia_gpu]
+        return [self.mock_other_gpu]
+      elif backend == 'tpu':
+        return [self.mock_tpu]
+      elif backend == 'cuda':
+        return [self.mock_nvidia_gpu]
+
+      raise AssertionError('Should not be called.')
+
+    self.mock_jax_devices.side_effect = devices_side_effect
+
+    expected_device, expected_backend_impl = expected
+    if expected_backend_impl == 'error':
+      with self.assertRaises(AssertionError):
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=input_device
+        )
+      return
+
+    actual_backend_impl, actual_device = (
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=input_device
+        )
+    )
+
+    self.assertEqual(actual_backend_impl, expected_backend_impl)
+    self.assertIsNotNone(actual_device)
+    self.assertEqual(actual_device.platform, expected_device)
+
+  @parameterized.named_parameters(
+      (f'{str(args[0])}_{str(args[1])}', *args)
+      for args in _DEFAULT_DEVICE_TEST_CASES
+  )
+  @mock.patch.dict(
+      os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'}
+  )
+  def test_resolve_without_device(
+      self,
+      default_device_str,
+      backend_impl_str,
+      expected,
+  ):
+    """Tests various combinations of jax.default_device and backend impls."""
+    default_devices = {
+        'cpu': [self.mock_cpu],
+        'gpu-nvidia': [self.mock_nvidia_gpu, self.mock_cpu],
+        'gpu-notnvidia': [self.mock_other_gpu, self.mock_cpu],
+        'tpu': [self.mock_tpu, self.mock_cpu],
+    }[default_device_str]
+
+    def devices_side_effect(backend=None):
+      if backend == 'cpu':
+        return [self.mock_cpu]  # CPU is always available
+      if backend == 'gpu' and default_device_str == 'gpu-notnvidia':
+        return [self.mock_other_gpu]
+      if backend == 'gpu' and default_device_str == 'gpu-nvidia':
+        return [self.mock_nvidia_gpu]
+      if backend == 'cuda' and default_device_str == 'gpu-nvidia':
+        return [self.mock_nvidia_gpu]
+      if backend == 'tpu' and default_device_str == 'tpu':
+        return [self.mock_tpu]
+      if backend is None:
+        return default_devices
+      if backend == 'cuda':
+        raise RuntimeError('cuda backend not supported')
+      raise AssertionError('jax.devices error')
+
+    self.mock_jax_devices.side_effect = devices_side_effect
+    default_device_side_effect_str = {
+        'cpu': 'cpu',
+        'gpu-nvidia': 'gpu',
+        'gpu-notnvidia': 'gpu',
+        'tpu': 'tpu',
+    }[default_device_str]
+    self.mock_default_backend.side_effect = (
+        lambda: default_device_side_effect_str
+    )
+
+    expected_device, expected_backend_impl = expected
+    if (
+        expected_backend_impl == 'error'
+        and default_device_str != 'gpu-nvidia'
+        and backend_impl_str == 'warp'
+    ):
+      with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=None
+        )
+      return
+
+    if expected_backend_impl == 'error':
+      with self.assertRaises(AssertionError):
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=None
+        )
+      return
+
+    actual_backend_impl, actual_device = (
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=None
+        )
+    )
+
+    self.assertEqual(actual_backend_impl, expected_backend_impl)
+    self.assertIsNotNone(actual_device)
+    self.assertEqual(actual_device.platform, expected_device)
+
+  @mock.patch.dict(os.environ, {'MJX_WARP_ENABLED': 'false'})
+  def test_resolve_warp_disabled(self):
+    """Tests behavior when MJX_WARP_ENABLED is false."""
+    self.mock_jax_devices.side_effect = lambda backend=None: (
+        [self.mock_nvidia_gpu, self.mock_cpu]
+        if backend is None
+        else ([self.mock_nvidia_gpu] if backend == 'gpu' else [self.mock_cpu])
+    )
+    self.mock_default_backend.side_effect = lambda: 'gpu'
+
+    # Default to JAX instead of WARP on NVIDIA GPU.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=None
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'gpu')
+
+    # Specifying an NVIDIA GPU should still choose JAX.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=self.mock_nvidia_gpu
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'gpu')
+
+    # Requesting warp explicitly should fail since it is disabled.
+    with self.assertRaises(AssertionError):
+      mjx_io._resolve_backend_impl_and_device(
+          backend_impl='warp', device=self.mock_nvidia_gpu
+      )
+    with self.assertRaises(AssertionError):
+      mjx_io._resolve_backend_impl_and_device(backend_impl='warp', device=None)
+
+  @mock.patch.dict(os.environ, {'MJX_C_DEFAULT_ENABLED': 'false'})
+  def test_resolve_c_disabled(self):
+    """Tests behavior when MJX_C_DEFAULT_ENABLED is false."""
+    # Users expect that CPU defaults to the JAX impl. But in the future, it will
+    # default to the C backend implementation. This test checks that
+    # MJX_C_DEFAULT_ENABLED=false defaults to the old behavior, until the
+    # migration to MJEP-15 is complete.
+    self.mock_jax_devices.side_effect = lambda backend=None: ([self.mock_cpu])
+    self.mock_default_backend.side_effect = lambda: 'cpu'
+
+    # Default to JAX instead of C on CPU.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=None
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'cpu')
+
+    # Specifing CPU should still choose JAX.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=self.mock_cpu
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'cpu')
+
+    # Specifying C should choose C!
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl='c', device=None
+    )
+    self.assertEqual(backend_impl, BackendImpl.C)
+    self.assertEqual(device.platform, 'cpu')
+
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl='c', device=self.mock_cpu
+    )
+    self.assertEqual(backend_impl, BackendImpl.C)
+    self.assertEqual(device.platform, 'cpu')
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/passive.py

@@ -23,9 +23,12 @@ from mujoco.mjx._src import scan
 from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
+from mujoco.mjx._src.types import OptionJAX
 # pylint: enable=g-importing-member
 
 
@@ -72,11 +75,11 @@ def _spring_damper(m: Model, d: Data) -> jax.Array:
   qfrc -= m.dof_damping * d.qvel
 
   # tendon-level spring-dampers
-  below, above = m.tendon_lengthspring.T - d.ten_length
+  below, above = m.tendon_lengthspring.T - d._impl.ten_length
   frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0)
   frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring)
-  frc_damper = -m.tendon_damping * d.ten_velocity
+  frc_damper = -m.tendon_damping * d._impl.ten_velocity
-  qfrc += d.ten_J.T @ (frc_spring + frc_damper)
+  qfrc += d._impl.ten_J.T @ (frc_spring + frc_damper)
 
   return qfrc
 
@@ -113,6 +116,9 @@ def _fluid(m: Model, d: Data) -> jax.Array:
 
 def passive(m: Model, d: Data) -> Data:
   """Adds all passive forces."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('passive requires JAX backend implementation.')
+
   if m.opt.disableflags & DisableBit.PASSIVE:
     return d.replace(qfrc_passive=jp.zeros(m.nv), qfrc_gravcomp=jp.zeros(m.nv))
 
@@ -124,7 +130,7 @@ def passive(m: Model, d: Data) -> Data:
     # add gravcomp unless added via actuators
     qfrc_passive += qfrc_gravcomp * (1 - m.jnt_actgravcomp[m.dof_jntid])
 
-  if m.opt.has_fluid_params:
+  if m.opt.has_fluid_params:  # pytype: disable=attribute-error
     qfrc_passive += _fluid(m, d)
 
   d = d.replace(qfrc_passive=qfrc_passive, qfrc_gravcomp=qfrc_gravcomp)

mjx/mujoco/mjx/_src/sensor.py

@@ -22,9 +22,12 @@ from mujoco.mjx._src import math
 from mujoco.mjx._src import ray
 from mujoco.mjx._src import smooth
 from mujoco.mjx._src import support
+from mujoco.mjx._src.types import BackendImpl
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import ObjType
 from mujoco.mjx._src.types import SensorType
 from mujoco.mjx._src.types import TrnType
@@ -32,7 +35,7 @@ from mujoco.mjx._src.types import TrnType
 import numpy as np
 
 
-def apply_cutoff(
+def _apply_cutoff(
     sensor: jax.Array, cutoff: jax.Array, data_type: int
 ) -> jax.Array:
   """Clip sensor to cutoff value."""
@@ -51,6 +54,8 @@ def apply_cutoff(
 
 def sensor_pos(m: Model, d: Data) -> Data:
   """Compute position-dependent sensors values."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('sensor_pos requires JAX backend implementation.')
 
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
@@ -163,15 +168,15 @@ def sensor_pos(m: Model, d: Data) -> Data:
         sensor, _ = jax.vmap(
             ray.ray, in_axes=(None, None, 0, 0, None, None, None)
         )(m, d, site_xpos, site_mat, (), True, sid)
-        sensors.append(apply_cutoff(sensor, cutoffs, data_type[0]))
+        sensors.append(_apply_cutoff(sensor, cutoffs, data_type[0]))
         adrs.append(adr[idxs])
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.JOINTPOS:
       sensor = d.qpos[m.jnt_qposadr[objid]]
     elif sensor_type == SensorType.TENDONPOS:
-      sensor = d.ten_length[objid]
+      sensor = d._impl.ten_length[objid]
     elif sensor_type == SensorType.ACTUATORPOS:
-      sensor = d.actuator_length[objid]
+      sensor = d._impl.actuator_length[objid]
     elif sensor_type == SensorType.BALLQUAT:
       jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
       quat = d.qpos[jnt_qposadr]
@@ -194,7 +199,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
         cutofft = cutoff[idxt]
         sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refidt)
         adrt = adr[idxt, None] + np.arange(3)[None]
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type in frame_axis:
@@ -214,7 +219,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
         cutofft = cutoff[idxt]
         sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refidt)
         adrt = adr[idxt, None] + np.arange(3)[None]
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.FRAMEQUAT:
@@ -255,7 +260,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
             )
         )(quat, refquat, refidt)
         adrt = adr[idxt, None] + np.arange(4)[None]
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.SUBTREECOM:
@@ -267,7 +272,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type
 
-    sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
+    sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
     adrs.append(adr)
 
   if not adrs:
@@ -282,6 +287,8 @@ def sensor_pos(m: Model, d: Data) -> Data:
 
 def sensor_vel(m: Model, d: Data) -> Data:
   """Compute velocity-dependent sensors values."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('sensor_vel requires JAX backend implementation.')
 
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
@@ -333,9 +340,9 @@ def sensor_vel(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.JOINTVEL:
       sensor = d.qvel[m.jnt_dofadr[objid]]
     elif sensor_type == SensorType.TENDONVEL:
-      sensor = d.ten_velocity[objid]
+      sensor = d._impl.ten_velocity[objid]
     elif sensor_type == SensorType.ACTUATORVEL:
-      sensor = d.actuator_velocity[objid]
+      sensor = d._impl.actuator_velocity[objid]
     elif sensor_type == SensorType.BALLANGVEL:
       jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
       sensor = d.qvel[jnt_dotadr]
@@ -393,20 +400,20 @@ def sensor_vel(m: Model, d: Data) -> Data:
 
         adrt = adr[idxt, None] + np.arange(3)[None]
 
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.SUBTREELINVEL:
-      sensor = d.subtree_linvel[objid]
+      sensor = d._impl.subtree_linvel[objid]
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     elif sensor_type == SensorType.SUBTREEANGMOM:
-      sensor = d.subtree_angmom[objid]
+      sensor = d._impl.subtree_angmom[objid]
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     else:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type
 
-    sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
+    sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
     adrs.append(adr)
 
   if not adrs:
@@ -421,6 +428,8 @@ def sensor_vel(m: Model, d: Data) -> Data:
 
 def sensor_acc(m: Model, d: Data) -> Data:
   """Compute acceleration/force-dependent sensors values."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('sensor_acc requires JAX backend implementation.')
 
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
@@ -460,14 +469,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
       # compute contact forces
       forces = []
       condim_ids = []
-      for dim in set(d.contact.dim):
+      for dim in set(d._impl.contact.dim):
         force, condim_id = support.contact_force_dim(m, d, dim)
         forces.append(force)
         condim_ids.append(condim_id)
       forces = jp.concatenate(forces)[np.argsort(np.concatenate(condim_ids))]
 
       # get bodies of contact geoms
-      conbody = jp.array(m.geom_bodyid)[d.contact.geom]
+      conbody = jp.array(m.geom_bodyid)[d._impl.contact.geom]
 
       # get site information
       site_bodyid = m.site_bodyid[objid]
@@ -477,12 +486,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
       site_type = m.site_type[objid]
       conbody0 = site_bodyid[:, None] == conbody[:, 0]
       conbody1 = site_bodyid[:, None] == conbody[:, 1]
-      contacts = (d.contact.efc_address >= 0)[None] & (conbody0 | conbody1)
+      contacts = (d._impl.contact.efc_address >= 0)[None] & (
+          conbody0 | conbody1
+      )
 
       # compute conray, flip if second body
       conray = jax.vmap(
           lambda frame, force: math.normalize(frame[0] * force[0])
-      )(d.contact.frame, forces)
+      )(d._impl.contact.frame, forces)
       conray = jp.where(conbody1[..., None], -conray, conray)
 
       # compute distance, mapping over sites and contacts
@@ -504,7 +515,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
             site_xpos[dist_id_site],
             site_xmat[dist_id_site],
             st,
-            d.contact.pos,
+            d._impl.contact.pos,
             conray[dist_id_site],
         )
         dist.append(jp.where(jp.isinf(dist_site), 0, dist_site))
@@ -526,14 +537,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
       bodyid = m.site_bodyid[objid]
       rot = d.site_xmat[objid]
       cvel = d.cvel[bodyid]
-      cacc = d.cacc[bodyid]
+      cacc = d._impl.cacc[bodyid]
       dif = d.site_xpos[objid] - d.subtree_com[m.body_rootid[bodyid]]
 
       sensor = _accelerometer(cvel, cacc, dif, rot)
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     elif sensor_type == SensorType.FORCE:
       bodyid = m.site_bodyid[objid]
-      cfrc_int = d.cfrc_int[bodyid]
+      cfrc_int = d._impl.cfrc_int[bodyid]
       site_xmat = d.site_xmat[objid]
       sensor = jax.vmap(lambda mat, vec: mat.T @ vec)(
           site_xmat, cfrc_int[:, 3:]
@@ -542,7 +553,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.TORQUE:
       bodyid = m.site_bodyid[objid]
       rootid = m.body_rootid[bodyid]
-      cfrc_int = d.cfrc_int[bodyid]
+      cfrc_int = d._impl.cfrc_int[bodyid]
       site_xmat = d.site_xmat[objid]
       dif = d.site_xpos[objid] - d.subtree_com[rootid]
       sensor = jax.vmap(
@@ -571,7 +582,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
         pos, bodyid = objtype_data[ot]
         pos = pos[objidt]
         bodyid = bodyid[objidt]
-        cacc = d.cacc[bodyid]
+        cacc = d._impl.cacc[bodyid]
 
         if sensor_type == SensorType.FRAMELINACC:
 
@@ -601,7 +612,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type
 
-    sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
+    sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
     adrs.append(adr)
 
   if not adrs:

mjx/mujoco/mjx/_src/sensor_test.py

@@ -64,14 +64,14 @@ class SensorTest(parameterized.TestCase):
     mujoco.mj_forward(m, d)
 
     mx = mjx.put_model(m)
-    dx = mjx.put_data(m, d).replace(
+    dx = mjx.put_data(m, d).tree_replace({
-        sensordata=jp.zeros_like(d.sensordata),
+        'sensordata': jp.zeros_like(d.sensordata),
-        subtree_linvel=jp.zeros_like(d.subtree_linvel),
+        '_impl.subtree_linvel': jp.zeros_like(d.subtree_linvel),
-        subtree_angmom=jp.zeros_like(d.subtree_angmom),
+        '_impl.subtree_angmom': jp.zeros_like(d.subtree_angmom),
-        cacc=jp.zeros_like(d.cacc),
+        '_impl.cacc': jp.zeros_like(d.cacc),
-        cfrc_int=jp.zeros_like(d.cfrc_int),
+        '_impl.cfrc_int': jp.zeros_like(d.cfrc_int),
-        cfrc_ext=jp.zeros_like(d.cfrc_ext),
+        '_impl.cfrc_ext': jp.zeros_like(d.cfrc_ext),
-    )
+    })
     dx = jax.jit(mjx.sensor_pos)(mx, dx)
     dx = jax.jit(mjx.sensor_vel)(mx, dx)
     dx = jax.jit(mjx.sensor_acc)(mx, dx)

mjx/mujoco/mjx/_src/smooth.py

@@ -23,10 +23,12 @@ from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import CamLightType
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import EqType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import TrnType
 from mujoco.mjx._src.types import WrapType
 # pylint: enable=g-importing-member
@@ -35,7 +37,6 @@ import numpy as np
 
 def kinematics(m: Model, d: Data) -> Data:
   """Converts position/velocity from generalized coordinates to maximal."""
-
   def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat):
     # calculate joint anchors, axes, body pos and quat in global frame
     # also normalize qpos while we're at it
@@ -131,6 +132,8 @@ def kinematics(m: Model, d: Data) -> Data:
 
 def com_pos(m: Model, d: Data) -> Data:
   """Maps inertias and motion dofs to global frame centered at subtree-CoM."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('com_pos requires JAX backend implementation.')
 
   # calculate center of mass of each subtree
   def subtree_sum(carry, xipos, body_mass):
@@ -162,7 +165,7 @@ def com_pos(m: Model, d: Data) -> Data:
   root_com = subtree_com[m.body_rootid]
   offset = d.xipos - root_com
   cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
-  d = d.replace(cinert=cinert)
+  d = d.tree_replace({'_impl.cinert': cinert})
 
   # map motion dofs to global frame centered at subtree_com
   def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis):
@@ -201,13 +204,16 @@ def com_pos(m: Model, d: Data) -> Data:
       d.xanchor,
       d.xaxis,
   )
-  d = d.replace(cdof=cdof)
+  d = d.tree_replace({'_impl.cdof': cdof})
 
   return d
 
 
 def camlight(m: Model, d: Data) -> Data:
   """Computes camera and light positions and orientations."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('camlight requires JAX backend implementation.')
+
   if m.ncam == 0:
     return d.replace(cam_xpos=jp.zeros((0, 3)), cam_xmat=jp.zeros((0, 3, 3)))
 
@@ -278,32 +284,33 @@ def camlight(m: Model, d: Data) -> Data:
 
 def crb(m: Model, d: Data) -> Data:
   """Runs composite rigid body inertia algorithm."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('crb requires JAX backend implementation.')
 
   def crb_fn(crb_child, crb_body):
     if crb_child is not None:
       crb_body += crb_child
     return crb_body
 
-  crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d.cinert, reverse=True)
+  crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d._impl.cinert, reverse=True)
   crb_body = crb_body.at[0].set(0.0)
-  d = d.replace(crb=crb_body)
+  d = d.tree_replace({'_impl.crb': crb_body})
 
   crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
-  crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
+  crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d._impl.cdof)
-  qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
+  qm = support.make_m(m, crb_cdof, d._impl.cdof, m.dof_armature)
-  d = d.replace(qM=qm)
+  d = d.tree_replace({'_impl.qM': qm})
-  if support.is_sparse(m) and d._qM_sparse.size > 0:  # pylint: disable=protected-access
-    d = d.replace(_qM_sparse=qm)
-
   return d
 
 
 def factor_m(m: Model, d: Data) -> Data:
   """Gets factorizaton of inertia-like matrix M, assumed spd."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('factor_m requires JAX backend implementation.')
 
   if not support.is_sparse(m):
-    qh, _ = jax.scipy.linalg.cho_factor(d.qM)
+    qh, _ = jax.scipy.linalg.cho_factor(d._impl.qM)
-    d = d.replace(qLD=qh)
+    d = d.tree_replace({'_impl.qLD': qh})
     return d
 
   # build up indices for where we will do backwards updates over qLD
@@ -325,7 +332,7 @@ def factor_m(m: Model, d: Data) -> Data:
           (out_beg, out_end, madr_d, madr_ij)
       )
 
-  qld = d.qM
+  qld = d._impl.qM
 
   for _, updates in sorted(updates.items(), reverse=True):
     # combine the updates into one update batch (per depth level)
@@ -353,20 +360,17 @@ def factor_m(m: Model, d: Data) -> Data:
   qld_diag = qld[m.dof_Madr]
   qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag)
 
-  d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag)
+  d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag})
-  if d._qLD_sparse.size > 0:  # pylint: disable=protected-access
-    d = d.replace(_qLD_sparse=d.qLD)
-  if d._qLDiagInv_sparse.size > 0:  # pylint: disable=protected-access
-    d = d.replace(_qLDiagInv_sparse=d.qLDiagInv)
-
   return d
 
 
 def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   """Computes sparse backsubstitution:  x = inv(L'*D*L)*y ."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('solve_m requires JAX backend implementation.')
 
   if not support.is_sparse(m):
-    return jax.scipy.linalg.cho_solve((d.qLD, False), x)
+    return jax.scipy.linalg.cho_solve((d._impl.qLD, False), x)
 
   depth = []
   for i in range(m.nv):
@@ -385,21 +389,23 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   # x <- inv(L') * x
   for _, vals in sorted(updates_j.items(), reverse=True):
     j, madr_ij, i = np.array(vals).T
-    x = x.at[j].add(-d.qLD[madr_ij] * x[i])
+    x = x.at[j].add(-d._impl.qLD[madr_ij] * x[i])
 
   # x <- inv(D) * x
-  x = x * d.qLDiagInv
+  x = x * d._impl.qLDiagInv
 
   # x <- inv(L) * x
   for _, vals in sorted(updates_i.items()):
     i, madr_ij, j = np.array(vals).T
-    x = x.at[i].add(-d.qLD[madr_ij] * x[j])
+    x = x.at[i].add(-d._impl.qLD[madr_ij] * x[j])
 
   return x
 
 
 def com_vel(m: Model, d: Data) -> Data:
   """Computes cvel, cdof_dot."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('com_vel requires JAX backend implementation.')
 
   # forward scan down tree: accumulate link center of mass velocity
   def fn(parent, jnt_typs, cdof, qvel):
@@ -431,17 +437,19 @@ def com_vel(m: Model, d: Data) -> Data:
       'jvv',
       'bv',
       m.jnt_type,
-      d.cdof,
+      d._impl.cdof,
       d.qvel,
   )
 
-  d = d.replace(cvel=cvel, cdof_dot=cdof_dot)
+  d = d.tree_replace({'cvel': cvel, '_impl.cdof_dot': cdof_dot})
 
   return d
 
 
 def subtree_vel(m: Model, d: Data) -> Data:
   """Subtree linear velocity and angular momentum."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('subtree_vel requires JAX backend implementation.')
 
   # bodywise quantities
   def _forward(cvel, xipos, ximat, subtree_com_root, mass, inertia):
@@ -529,7 +537,10 @@ def subtree_vel(m: Model, d: Data) -> Data:
       reverse=True,
   )
 
-  return d.replace(subtree_linvel=subtree_linvel, subtree_angmom=subtree_angmom)
+  return d.tree_replace({
+      '_impl.subtree_linvel': subtree_linvel,
+      '_impl.subtree_angmom': subtree_angmom,
+  })
 
 
 def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
@@ -537,6 +548,8 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
 
   flg_acc=False removes inertial term.
   """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('rne requires JAX backend implementation.')
 
   # forward scan over tree: accumulate link center of mass acceleration
   def cacc_fn(cacc, cdof_dot, qvel, cdof, qacc):
@@ -555,7 +568,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
     return cacc
 
   cacc = scan.body_tree(
-      m, cacc_fn, 'vvvv', 'b', d.cdof_dot, d.qvel, d.cdof, d.qacc
+      m, cacc_fn, 'vvvv', 'b', d._impl.cdof_dot, d.qvel, d._impl.cdof, d.qacc
   )
 
   def frc(cinert, cacc, cvel):
@@ -564,7 +577,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
 
     return frc
 
-  loc_cfrc = jax.vmap(frc)(d.cinert, cacc, d.cvel)
+  loc_cfrc = jax.vmap(frc)(d._impl.cinert, cacc, d.cvel)
 
   # backward scan up tree: accumulate body forces
   def cfrc_fn(cfrc_child, cfrc):
@@ -573,7 +586,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
     return cfrc
 
   cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True)
-  qfrc_bias = jax.vmap(jp.dot)(d.cdof, cfrc[jp.array(m.dof_bodyid)])
+  qfrc_bias = jax.vmap(jp.dot)(d._impl.cdof, cfrc[jp.array(m.dof_bodyid)])
 
   d = d.replace(qfrc_bias=qfrc_bias)
 
@@ -582,6 +595,8 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
 
 def rne_postconstraint(m: Model, d: Data) -> Data:
   """RNE with complete data: compute cacc, cfrc_ext, cfrc_int."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('rne_postconstraint requires JAX backend implementation.')
 
   def _transform_force(frc, offset):
     force, torque = jp.split(frc, 2)
@@ -602,7 +617,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
   # compute contact forces for each condim
   forces = []
   condim_idx = []
-  for dim in set(d.contact.dim):
+  for dim in set(d._impl.contact.dim):
     force, idx = support.contact_force_dim(m, d, dim)
     forces.append(force)
     condim_idx.append(idx)
@@ -629,10 +644,10 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
       )
 
     condim_idx = jp.concatenate(condim_idx)
-    frame = d.contact.frame[condim_idx]
+    frame = d._impl.contact.frame[condim_idx]
-    pos = d.contact.pos[condim_idx]
+    pos = d._impl.contact.pos[condim_idx]
-    id1 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 0]]
+    id1 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 0]]
-    id2 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 1]]
+    id2 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 1]]
     com1 = d.subtree_com[jp.array(m.body_rootid)][id1]
     com2 = d.subtree_com[jp.array(m.body_rootid)][id2]
 
@@ -668,8 +683,8 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
     )
 
     # cacc = cacc_parent + cdofdot * qvel + cdof * qacc
-    cacc_vel = d.cdof_dot.T @ (mask * d.qvel)
+    cacc_vel = d._impl.cdof_dot.T @ (mask * d.qvel)
-    cacc_acc = d.cdof.T @ (mask * d.qacc)
+    cacc_acc = d._impl.cdof.T @ (mask * d.qacc)
     cacc = cacc_parent + cacc_vel + cacc_acc
 
     # cfrc_body = cinert * cacc + cvel x (cinert * cvel)
@@ -687,7 +702,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
       'bbbbb',
       'bb',
       cfrc_ext,
-      d.cinert,
+      d._impl.cinert,
       d.cvel,
       jp.array(m.body_dofadr),
       jp.array(m.body_dofnum),
@@ -704,11 +719,18 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
   )
 
   # update data
-  return d.replace(cacc=cacc, cfrc_int=cfrc_int, cfrc_ext=cfrc_ext)
+  return d.tree_replace({
+      '_impl.cacc': cacc,
+      '_impl.cfrc_int': cfrc_int,
+      '_impl.cfrc_ext': cfrc_ext,
+  })
 
 
 def tendon(m: Model, d: Data) -> Data:
   """Computes tendon lengths and moments."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('tendon requires JAX backend implementation.')
+
   if not m.ntendon:
     return d
 
@@ -836,8 +858,8 @@ def tendon(m: Model, d: Data) -> Data:
   # wrap inside
   # TODO(taylorhowell): check that is_wrap_inside is consistent with
   # site and geom relative positions
-  (wrap_inside_id,) = np.nonzero(m.is_wrap_inside)
+  (wrap_inside_id,) = np.nonzero(m._impl.is_wrap_inside)
-  (wrap_outside_id,) = np.nonzero(~m.is_wrap_inside)
+  (wrap_outside_id,) = np.nonzero(~m._impl.is_wrap_inside)
 
   # compute geom wrap length and connect points (if wrap occurs)
   v_wrap = jax.vmap(
@@ -853,9 +875,9 @@ def tendon(m: Model, d: Data) -> Data:
       has_sidesite[wrap_inside_id],
       is_sphere[wrap_inside_id],
       True,
-      m.wrap_inside_maxiter,
+      m._impl.wrap_inside_maxiter,
-      m.wrap_inside_tolerance,
+      m._impl.wrap_inside_tolerance,
-      m.wrap_inside_z_init,
+      m._impl.wrap_inside_z_init,
   )
 
   lengths_outside, pnt0_outside, pnt1_outside = v_wrap(
@@ -868,9 +890,9 @@ def tendon(m: Model, d: Data) -> Data:
       has_sidesite[wrap_outside_id],
       is_sphere[wrap_outside_id],
       False,
-      m.wrap_inside_maxiter,
+      m._impl.wrap_inside_maxiter,
-      m.wrap_inside_tolerance,
+      m._impl.wrap_inside_tolerance,
-      m.wrap_inside_z_init,
+      m._impl.wrap_inside_z_init,
   )
 
   wrap_id = np.argsort(np.concatenate([wrap_inside_id, wrap_outside_id]))
@@ -952,12 +974,14 @@ def tendon(m: Model, d: Data) -> Data:
   )
 
   # assemble length and moment
-  ten_length = jp.zeros_like(d.ten_length).at[tendon_id_jnt].set(length_jnt)
+  ten_length = (
+      jp.zeros_like(d._impl.ten_length).at[tendon_id_jnt].set(length_jnt)
+  )
   ten_length = ten_length.at[tendon_id_site].add(length_site)
   ten_length = ten_length.at[tendon_id_geom].add(length_geom)
 
   ten_moment = (
-      jp.zeros_like(d.ten_J)
+      jp.zeros_like(d._impl.ten_J)
       .at[adr_moment_jnt, dofadr_moment_jnt]
       .set(moment_jnt)
   )
@@ -1020,14 +1044,14 @@ def tendon(m: Model, d: Data) -> Data:
       [wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)]
   ).reshape((m.nwrap, 2))
 
-  return d.replace(
+  return d.tree_replace({
-      ten_length=ten_length,
+      '_impl.ten_length': ten_length,
-      ten_J=ten_moment,
+      '_impl.ten_J': ten_moment,
-      ten_wrapadr=jp.array(ten_wrapadr, dtype=int),
+      '_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int),
-      ten_wrapnum=jp.array(ten_wrapnum, dtype=int),
+      '_impl.ten_wrapnum': jp.array(ten_wrapnum, dtype=int),
-      wrap_xpos=wrap_xpos,
+      '_impl.wrap_xpos': wrap_xpos,
-      wrap_obj=jp.array(wrap_obj, dtype=int),
+      '_impl.wrap_obj': jp.array(wrap_obj, dtype=int),
-  )
+  })
 
 
 def _site_dof_mask(m: Model) -> np.ndarray:
@@ -1061,6 +1085,9 @@ def _site_dof_mask(m: Model) -> np.ndarray:
 
 def transmission(m: Model, d: Data) -> Data:
   """Computes actuator/transmission lengths and moments."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('transmission requires JAX backend implementation.')
+
   if not m.nu:
     return d
 
@@ -1120,8 +1147,8 @@ def transmission(m: Model, d: Data) -> Data:
       wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:]))
       moment = jac @ wrench
     elif trntype == TrnType.TENDON:
-      length = d.ten_length[trnid[0]] * gear[:1]
+      length = d._impl.ten_length[trnid[0]] * gear[:1]
-      moment = d.ten_J[trnid[0]] * gear[0]
+      moment = d._impl.ten_J[trnid[0]] * gear[0]
     else:
       raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
 
@@ -1153,5 +1180,7 @@ def transmission(m: Model, d: Data) -> Data:
   length = length.reshape((m.nu,))
   moment = moment.reshape((m.nu, m.nv))
 
-  d = d.replace(actuator_length=length, actuator_moment=moment)
+  d = d.tree_replace(
+      {'_impl.actuator_length': length, '_impl.actuator_moment': moment}
+  )
   return d

mjx/mujoco/mjx/_src/smooth_test.py

@@ -20,7 +20,7 @@ import jax
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
-from mujoco.mjx._src.types import ConeType
+from mujoco.mjx._src.types import ConeType  # pylint: disable=g-importing-member
 import numpy as np
 
 # tolerance for difference between MuJoCo and MJX smooth calculations - mostly
@@ -78,30 +78,27 @@ class SmoothTest(absltest.TestCase):
     # com_pos
     dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'subtree_com')
-    _assert_attr_eq(d, dx, 'cinert')
+    _assert_attr_eq(d, dx._impl, 'cinert')
-    _assert_attr_eq(d, dx, 'cdof')
+    _assert_attr_eq(d, dx._impl, 'cdof')
     # camlight
     dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'cam_xpos')
     _assert_eq(d.cam_xmat.reshape((-1, 3, 3)), dx.cam_xmat, 'cam_xmat')
     # crb
     dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
-    _assert_attr_eq(d, dx, 'crb')
+    _assert_attr_eq(d, dx._impl, 'crb')
-    _assert_attr_eq(d, dx, 'qM')
+    _assert_attr_eq(d, dx._impl, 'qM')
-    _assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
     # factor_m
     dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
     qLDLegacy = np.zeros(mx.nM)  # pylint:disable=invalid-name
     for i in range(m.nM):
       qLDLegacy[d.mapM2M[i]] = d.qLD[i]
-    _assert_eq(qLDLegacy, dx.qLD, 'qLD')
+    _assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
-    _assert_attr_eq(d, dx, 'qLDiagInv')
+    _assert_attr_eq(d, dx._impl, 'qLDiagInv')
-    _assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
-    _assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')
     # com_vel
     dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'cvel')
-    _assert_attr_eq(d, dx, 'cdof_dot')
+    _assert_attr_eq(d, dx._impl, 'cdof_dot')
     # rne
     dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'qfrc_bias')
@@ -122,11 +119,11 @@ class SmoothTest(absltest.TestCase):
     mujoco.mj_forward(m, d)
     # tendon
     dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d))
-    _assert_attr_eq(d, dx, 'ten_J')
+    _assert_attr_eq(d, dx._impl, 'ten_J')
-    _assert_attr_eq(d, dx, 'ten_length')
+    _assert_attr_eq(d, dx._impl, 'ten_length')
     # transmission
     dx = jax.jit(mjx.transmission)(mx, dx)
-    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx._impl, 'actuator_length')
 
     # convert sparse actuator_moment to dense representation
     moment = np.zeros((m.nu, m.nv))
@@ -137,7 +134,7 @@ class SmoothTest(absltest.TestCase):
         d.moment_rowadr,
         d.moment_colind,
     )
-    _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
+    _assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment')
 
   def test_disable_gravity(self):
     m = mujoco.MjModel.from_xml_string("""
@@ -197,7 +194,7 @@ class SmoothTest(absltest.TestCase):
 
     mujoco.mj_transmission(m, d)
     dx = jax.jit(mjx.transmission)(mx, dx)
-    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx._impl, 'actuator_length')
 
     # convert sparse actuator_moment to dense representation
     moment = np.zeros((m.nu, m.nv))
@@ -208,7 +205,7 @@ class SmoothTest(absltest.TestCase):
         d.moment_rowadr,
         d.moment_colind,
     )
-    _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
+    _assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment')
 
   def test_subtree_vel(self):
     """Tests MJX subtree_vel function matches MuJoCo mj_subtreeVel."""
@@ -226,8 +223,8 @@ class SmoothTest(absltest.TestCase):
     mujoco.mj_subtreeVel(m, d)
     dx = jax.jit(mjx.subtree_vel)(mx, dx)
 
-    _assert_attr_eq(d, dx, 'subtree_linvel')
+    _assert_attr_eq(d, dx._impl, 'subtree_linvel')
-    _assert_attr_eq(d, dx, 'subtree_angmom')
+    _assert_attr_eq(d, dx._impl, 'subtree_angmom')
 
 
 class RnePostConstraintTest(parameterized.TestCase):
@@ -278,9 +275,9 @@ class RnePostConstraintTest(parameterized.TestCase):
     mujoco.mj_rnePostConstraint(m, d)
     dx = jax.jit(mjx.rne_postconstraint)(mx, dx)
 
-    _assert_eq(d.cacc, dx.cacc, 'cacc')
+    _assert_eq(d.cacc, dx._impl.cacc, 'cacc')
-    _assert_eq(d.cfrc_ext, dx.cfrc_ext, 'cfrc_ext')
+    _assert_eq(d.cfrc_ext, dx._impl.cfrc_ext, 'cfrc_ext')
-    _assert_eq(d.cfrc_int, dx.cfrc_int, 'cfrc_int')
+    _assert_eq(d.cfrc_int, dx._impl.cfrc_int, 'cfrc_int')
 
 
 class TendonTest(parameterized.TestCase):
@@ -312,12 +309,12 @@ class TendonTest(parameterized.TestCase):
     mujoco.mj_forward(m, d)
     dx = jax.jit(mjx.forward)(mx, dx)
 
-    _assert_eq(d.ten_length, dx.ten_length, 'ten_length')
+    _assert_eq(d.ten_length, dx._impl.ten_length, 'ten_length')
-    _assert_eq(d.ten_J, dx.ten_J, 'ten_J')
+    _assert_eq(d.ten_J, dx._impl.ten_J, 'ten_J')
-    _assert_eq(d.ten_wrapnum, dx.ten_wrapnum, 'ten_wrapnum')
+    _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')
-    _assert_eq(d.ten_wrapadr, dx.ten_wrapadr, 'ten_wrapadr')
+    _assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr')
-    _assert_eq(d.wrap_obj, dx.wrap_obj, 'wrap_obj')
+    _assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj')
-    _assert_eq(d.wrap_xpos, dx.wrap_xpos, 'wrap_xpos')
+    _assert_eq(d.wrap_xpos, dx._impl.wrap_xpos, 'wrap_xpos')
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/solver.py

@@ -24,8 +24,10 @@ from mujoco.mjx._src import support
 from mujoco.mjx._src.dataclasses import PyTreeNode
 from mujoco.mjx._src.types import ConeType
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import SolverType
 # pylint: enable=g-importing-member
 
@@ -73,14 +75,19 @@ class Context(PyTreeNode):
 
   @classmethod
   def create(cls, m: Model, d: Data, grad: bool = True) -> 'Context':
-    jaref = d.efc_J @ d.qacc - d.efc_aref
+    if not isinstance(d._impl, DataJAX):
+      raise ValueError(
+          'Constraint context requires JAX backend implementation.'
+      )
+
+    jaref = d._impl.efc_J @ d.qacc - d._impl.efc_aref
     # TODO(robotics-team): determine nv at which sparse mul is faster
     ma = support.mul_m(m, d, d.qacc)
     nv_0 = jp.zeros(m.nv)
     fri = 0.0
     if m.opt.cone == ConeType.ELLIPTIC:
-      friction = d.contact.friction[d.contact.dim > 1]
+      friction = d._impl.contact.friction[d._impl.contact.dim > 1]
-      dim = d.contact.dim[d.contact.dim > 1]
+      dim = d._impl.contact.dim[d._impl.contact.dim > 1]
       mu = friction[:, 0] / jp.sqrt(m.opt.impratio)
       fri = jp.concatenate((mu[:, None], friction), axis=1)
       for condim in (3, 4, 6):
@@ -90,7 +97,7 @@ class Context(PyTreeNode):
         qacc=d.qacc,
         qfrc_constraint=d.qfrc_constraint,
         Jaref=jaref,
-        efc_force=d.efc_force,
+        efc_force=d._impl.efc_force,
         Ma=ma,
         grad=nv_0,
         Mgrad=nv_0,
@@ -145,18 +152,20 @@ class _LSPoint(PyTreeNode):
   ) -> '_LSPoint':
     """Creates a linesearch point with first and second derivatives."""
     # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
+    if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+      raise ValueError('LSPoint requires JAX backend implementation.')
 
     cost, deriv_0, deriv_1 = 0.0, 0.0, 0.0
     quad_total = quad_gauss
     x = ctx.Jaref + alpha * jv
-    active = (x < 0).at[: d.ne + d.nf].set(True)
+    active = (x < 0).at[: d._impl.ne + d._impl.nf].set(True)
 
-    dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
+    dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss
     if (dof_fl.any() or ten_fl.any()) and not (
         m.opt.disableflags & DisableBit.FRICTIONLOSS
     ):
-      f = d.efc_frictionloss
+      f = d._impl.efc_frictionloss
-      r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
+      r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL)
       rf, z = r * f, jp.zeros_like(f)
       linear_neg = (x <= -rf)[:, None]
       linear_pos = (x >= rf)[:, None]
@@ -174,13 +183,13 @@ class _LSPoint(PyTreeNode):
       middle_zone = (tsqr > 0) & (n < (mu * t)) & ((mu * n + t) > 0)
 
       # quadratic cost for equality, friction, limits, frictionless contacts
-      dim1 = d.contact.efc_address[d.contact.dim == 1]
+      dim1 = d._impl.contact.efc_address[d._impl.contact.dim == 1]
-      nefl = d.ne + d.nf + d.nl
+      nefl = d._impl.ne + d._impl.nf + d._impl.nl
       active = active.at[nefl:].set(False).at[dim1].set(active[dim1])
       quad_efld = jax.vmap(jp.multiply)(quad, active)
       quad_total += jp.sum(quad_efld, axis=0)
       # elliptic bottom zone: quadratic cost
-      efc_elliptic = d.contact.efc_address[d.contact.dim > 1]
+      efc_elliptic = d._impl.contact.efc_address[d._impl.contact.dim > 1]
       quad_c = jax.vmap(jp.multiply)(quad[efc_elliptic], bottom_zone)
       quad_total += jp.sum(quad_c, axis=0)
       # elliptic middle zone
@@ -254,17 +263,20 @@ def _update_constraint(m: Model, d: Data, ctx: Context) -> Context:
   Returns:
     context with new constraint force and costs
   """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_update_constraint requires JAX backend implementation.')
+
   # ne constraints are always active, nf are conditionally active, others are
   # non-negative constraints.
-  active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True)
+  active = (ctx.Jaref < 0).at[: d._impl.ne + d._impl.nf].set(True)
 
-  floss_force, floss_cost = jp.zeros(d.nefc), 0.0
+  floss_force, floss_cost = jp.zeros(d._impl.nefc), 0.0
-  dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
+  dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss
   if (dof_fl.any() or ten_fl.any()) and not (
       m.opt.disableflags & DisableBit.FRICTIONLOSS
   ):
-    f = d.efc_frictionloss
+    f = d._impl.efc_frictionloss
-    r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
+    r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL)
     linear_neg = (ctx.Jaref <= -r * f) * (f > 0)
     linear_pos = (ctx.Jaref >= r * f) * (f > 0)
     active = active & ~linear_neg & ~linear_pos
@@ -274,13 +286,13 @@ def _update_constraint(m: Model, d: Data, ctx: Context) -> Context:
     floss_cost = floss_cost.sum()
 
   if m.opt.cone == ConeType.PYRAMIDAL:
-    efc_force = d.efc_D * -ctx.Jaref * active + floss_force
+    efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force
-    cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
+    cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active)
     dm, u, h = 0.0, 0.0, 0.0
   elif m.opt.cone == ConeType.ELLIPTIC:
-    friction = d.contact.friction[d.contact.dim > 1]
+    friction = d._impl.contact.friction[d._impl.contact.dim > 1]
-    efc_address = d.contact.efc_address[d.contact.dim > 1]
+    efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1]
-    dim = d.contact.dim[d.contact.dim > 1]
+    dim = d._impl.contact.dim[d._impl.contact.dim > 1]
     # to prevent out of range append zeros to ctx.Jaref
     slice_fn = jax.vmap(
         lambda x: jax.lax.dynamic_slice(
@@ -297,12 +309,12 @@ def _update_constraint(m: Model, d: Data, ctx: Context) -> Context:
       adr_i.extend(range(addr, addr + condim))
       adr_j.extend([i] * condim)
     active = active.at[jp.array(adr_i)].set(bottom_zone[jp.array(adr_j)])
-    efc_force = d.efc_D * -ctx.Jaref * active + floss_force
+    efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force
-    cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
+    cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active)
 
     # middle zone: cone
     middle_zone = (t > 0) & (n < (mu * t)) & ((mu * n + t) > 0)
-    dm = d.efc_D[efc_address] / jp.maximum(
+    dm = d._impl.efc_D[efc_address] / jp.maximum(
         mu * mu * (1 + mu * mu), mujoco.mjMINVAL
     )
     nmt = n - mu * t
@@ -339,7 +351,7 @@ def _update_constraint(m: Model, d: Data, ctx: Context) -> Context:
   else:
     raise NotImplementedError(f'unsupported cone type: {m.opt.cone}')
 
-  qfrc_constraint = d.efc_J.T @ efc_force
+  qfrc_constraint = d._impl.efc_J.T @ efc_force
   gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth)
   ctx = ctx.replace(
       qfrc_constraint=qfrc_constraint,
@@ -371,6 +383,8 @@ def _update_gradient(m: Model, d: Data, ctx: Context) -> Context:
   Raises:
     NotImplementedError: for unsupported solver type
   """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_update_gradient requires JAX backend implementation.')
 
   grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint
 
@@ -378,16 +392,16 @@ def _update_gradient(m: Model, d: Data, ctx: Context) -> Context:
     mgrad = smooth.solve_m(m, d, grad)
   elif m.opt.solver == SolverType.NEWTON:
     if m.opt.cone == ConeType.ELLIPTIC:
-      cm = jp.diag(d.efc_D * ctx.active)
+      cm = jp.diag(d._impl.efc_D * ctx.active)
-      efc_address = d.contact.efc_address[d.contact.dim > 1]
+      efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1]
-      dim = d.contact.dim[d.contact.dim > 1]
+      dim = d._impl.contact.dim[d._impl.contact.dim > 1]
       # set efc of cone H along diagonal
       for i, (condim, addr) in enumerate(zip(dim, efc_address)):
         h_cone = ctx.h[i, :condim, :condim]
         cm = cm.at[addr : addr + condim, addr : addr + condim].add(h_cone)
-      h = d.efc_J.T @ cm @ d.efc_J
+      h = d._impl.efc_J.T @ cm @ d._impl.efc_J
     else:
-      h = (d.efc_J.T * d.efc_D * ctx.active) @ d.efc_J
+      h = (d._impl.efc_J.T * d._impl.efc_D * ctx.active) @ d._impl.efc_J
     h = support.full_m(m, d) + h
     h_ = jax.scipy.linalg.cho_factor(h)
     mgrad = jax.scipy.linalg.cho_solve(h_, grad)
@@ -414,12 +428,15 @@ def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
   Returns:
     updated context with new qacc, Ma, Jaref
   """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_lineasearch requires JAX backend implementation.')
+
   smag = math.norm(ctx.search) * m.stat.meaninertia * max(1, m.nv)
   gtol = m.opt.tolerance * m.opt.ls_tolerance * smag
 
   # compute Mv, Jv
   mv = support.mul_m(m, d, ctx.search)
-  jv = d.efc_J @ ctx.search
+  jv = d._impl.efc_J @ ctx.search
 
   # prepare quadratics
   quad_gauss = jp.stack((
@@ -428,13 +445,15 @@ def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
       0.5 * jp.dot(ctx.search, mv),
   ))
   quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
-  quad = (quad * d.efc_D).T
+  quad = (quad * d._impl.efc_D).T
   uu, v0, uv, vv = 0.0, 0.0, 0.0, 0.0
   if m.opt.cone == ConeType.ELLIPTIC:
-    mask = d.contact.dim > 1
+    mask = d._impl.contact.dim > 1
     # complete vector quadratic (for bottom zone)
     efc_con, efc_fri = [], []
-    for condim, addr in zip(d.contact.dim[mask], d.contact.efc_address[mask]):
+    for condim, addr in zip(
+        d._impl.contact.dim[mask], d._impl.contact.efc_address[mask]
+    ):
       efc_con.extend([addr] * (condim - 1))
       efc_fri.extend(range(addr + 1, addr + condim))
     quad = quad.at[jp.array(efc_con)].add(quad[jp.array(efc_fri)])
@@ -446,7 +465,7 @@ def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
             jp.concatenate((jv, jp.zeros(3))), (x,), (6,)
         )
     )
-    efc_elliptic = d.contact.efc_address[mask]
+    efc_elliptic = d._impl.contact.efc_address[mask]
     v = jv_fn(efc_elliptic) * ctx.fri
     uu = jp.sum(ctx.u[:, 1:] * ctx.u[:, 1:], axis=1)
     v0 = v[:, 0]
@@ -571,11 +590,11 @@ def solve(m: Model, d: Data) -> Data:
   else:
     ctx = jax.lax.while_loop(cond, body, ctx)
 
-  d = d.replace(
+  d = d.tree_replace({
-      qacc_warmstart=ctx.qacc,
+      'qacc_warmstart': ctx.qacc,
-      qacc=ctx.qacc,
+      'qfrc_constraint': ctx.qfrc_constraint,
-      qfrc_constraint=ctx.qfrc_constraint,
+      'qacc': ctx.qacc,
-      efc_force=ctx.efc_force,
+      '_impl.efc_force': ctx.efc_force,
-  )
+  })
 
   return d

mjx/mujoco/mjx/_src/solver_test.py

@@ -87,8 +87,8 @@ class SolverTest(parameterized.TestCase):
 
       # MJX finds very similar solutions with the newton solver
       if solver_ == mujoco.mjtSolver.mjSOL_NEWTON:
-        nnz = dx.efc_J.any(axis=1)
+        nnz = dx._impl.efc_J.any(axis=1)
-        _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+        _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
         _assert_attr_eq(d, dx, 'qfrc_constraint')
         _assert_attr_eq(d, dx, 'qacc')
 
@@ -108,9 +108,9 @@ class SolverTest(parameterized.TestCase):
     mujoco.mj_forward(m, d)
     mx = mjx.put_model(m)
     dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
-    nnz = dx.efc_J.any(axis=1)
+    nnz = dx._impl.efc_J.any(axis=1)
     # even without warmstart, newton converges quickly
-    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-4)
+    _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force', tol=2e-4)
 
   def test_sparse(self):
     """Test solver works with sparse mass matrices."""
@@ -130,8 +130,8 @@ class SolverTest(parameterized.TestCase):
 
     _assert_attr_eq(d, dx, 'qacc')
     _assert_attr_eq(d, dx, 'qfrc_constraint')
-    nnz = dx.efc_J.any(axis=1)
+    nnz = dx._impl.efc_J.any(axis=1)
-    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+    _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
 
   def test_quad_frictionloss(self):
     """Test a case with quadratic frictionloss constraints."""
@@ -144,8 +144,8 @@ class SolverTest(parameterized.TestCase):
 
     _assert_attr_eq(d, dx, 'qacc')
     _assert_attr_eq(d, dx, 'qfrc_constraint')
-    nnz = dx.efc_J.any(axis=1)
+    nnz = dx._impl.efc_J.any(axis=1)
-    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+    _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
 
   # TODO(taylorhowell): condim=1 with ConeType.ELLIPTIC
   @parameterized.product(condim=(3, 4, 6), cone=tuple(ConeType))

mjx/mujoco/mjx/_src/support.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """Engine support functions."""
+
 from collections.abc import Iterable, Sequence
 from typing import Optional, Tuple, Union
 
@@ -92,7 +93,7 @@ def full_m(m: Model, d: Data) -> jax.Array:
   """Reconstitute dense mass matrix from qM."""
 
   if not is_sparse(m):
-    return d.qM
+    return d._impl.qM  # pytype: disable=attribute-error
 
   ij = []
   for i in range(m.nv):
@@ -103,7 +104,7 @@ def full_m(m: Model, d: Data) -> jax.Array:
 
   i, j = (jp.array(x) for x in zip(*ij))
 
-  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM)
+  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.qM)  # pytype: disable=attribute-error
 
   # also set upper triangular
   mat = mat + jp.tril(mat, -1).T
@@ -115,9 +116,9 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
   """Multiply vector by inertia matrix."""
 
   if not is_sparse(m):
-    return d.qM @ vec
+    return d._impl.qM @ vec  # pytype: disable=attribute-error
 
-  diag_mul = d.qM[jp.array(m.dof_Madr)] * vec
+  diag_mul = d._impl.qM[jp.array(m.dof_Madr)] * vec  # pytype: disable=attribute-error
 
   is_, js, madr_ijs = [], [], []
   for i in range(m.nv):
@@ -131,8 +132,8 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
 
   i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
 
-  out = diag_mul.at[i].add(d.qM[madr_ij] * vec[j])
+  out = diag_mul.at[i].add(d._impl.qM[madr_ij] * vec[j])  # pytype: disable=attribute-error
-  out = out.at[j].add(d.qM[madr_ij] * vec[i])
+  out = out.at[j].add(d._impl.qM[madr_ij] * vec[i])  # pytype: disable=attribute-error
 
   return out
 
@@ -147,9 +148,9 @@ def jac(
   mask = mask[jp.array(m.dof_bodyid)] > 0
 
   offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
-  jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d.cdof)
+  jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d._impl.cdof)  # pytype: disable=attribute-error
   jacp = jax.vmap(jp.multiply)(jacp, mask)
-  jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask)
+  jacr = jax.vmap(jp.multiply)(d._impl.cdof[:, :3], mask)  # pytype: disable=attribute-error
 
   return jacp, jacr
 
@@ -548,20 +549,20 @@ def contact_force(
     m: Model, d: Data, contact_id: int, to_world_frame: bool = False
 ) -> jax.Array:
   """Extract 6D force:torque for one contact, in contact frame by default."""
-  efc_address = d.contact.efc_address[contact_id]
+  efc_address = d._impl.contact.efc_address[contact_id]  # pytype: disable=attribute-error
-  condim = d.contact.dim[contact_id]
+  condim = d._impl.contact.dim[contact_id]  # pytype: disable=attribute-error
   if m.opt.cone == ConeType.PYRAMIDAL:
     force = _decode_pyramid(
-        d.efc_force[efc_address:], d.contact.friction[contact_id], condim
+        d._impl.efc_force[efc_address:], d._impl.contact.friction[contact_id], condim  # pytype: disable=attribute-error
     )
   elif m.opt.cone == ConeType.ELLIPTIC:
-    force = d.efc_force[efc_address : efc_address + condim]
+    force = d._impl.efc_force[efc_address : efc_address + condim]  # pytype: disable=attribute-error
     force = jp.concatenate([force, jp.zeros((6 - condim))])
   else:
     raise ValueError(f'Unknown cone type: {m.opt.cone}')
 
   if to_world_frame:
-    force = force.reshape((-1, 3)) @ d.contact.frame[contact_id]
+    force = force.reshape((-1, 3)) @ d._impl.contact.frame[contact_id]  # pytype: disable=attribute-error
     force = force.reshape(-1)
 
   return force * (efc_address >= 0)
@@ -572,21 +573,21 @@ def contact_force_dim(
 ) -> Tuple[jax.Array, np.ndarray]:
   """Extract 6D force:torque for contacts with dimension dim."""
   # valid contact and condim indices
-  idx_dim = (d.contact.efc_address >= 0) & (d.contact.dim == dim)
+  idx_dim = (d._impl.contact.efc_address >= 0) & (d._impl.contact.dim == dim)  # pytype: disable=attribute-error
 
   # contact force from efc
   if m.opt.cone == ConeType.PYRAMIDAL:
     efc_address = (
-        d.contact.efc_address[idx_dim, None]
+        d._impl.contact.efc_address[idx_dim, None]  # pytype: disable=attribute-error
         + np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None]
     )
-    efc_force = d.efc_force[efc_address]
+    efc_force = d._impl.efc_force[efc_address]  # pytype: disable=attribute-error
     force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))(
-        efc_force, d.contact.friction[idx_dim], dim
+        efc_force, d._impl.contact.friction[idx_dim], dim  # pytype: disable=attribute-error
     )
   elif m.opt.cone == ConeType.ELLIPTIC:
-    efc_address = d.contact.efc_address[idx_dim, None] + np.arange(dim)[None]
+    efc_address = d._impl.contact.efc_address[idx_dim, None] + np.arange(dim)[None]  # pytype: disable=attribute-error
-    force = d.efc_force[efc_address]
+    force = d._impl.efc_force[efc_address]  # pytype: disable=attribute-error
     force = jp.hstack([force, jp.zeros((force.shape[0], 6 - dim))])
   else:
     raise ValueError(f'Unknown cone type: {m.opt.cone}.')

mjx/mujoco/mjx/_src/support_test.py

@@ -385,7 +385,7 @@ class SupportTest(parameterized.TestCase):
 
     # map MJX contacts to MJ ones
     def _find(g):
-      val = (g == dx.contact.geom).sum(axis=1)
+      val = (g == dx._impl.contact.geom).sum(axis=1)
       return np.where(val == 2)[0][0]
 
     contact_id_map = {i: _find(d.contact.geom[i]) for i in range(d.ncon)}
@@ -399,7 +399,7 @@ class SupportTest(parameterized.TestCase):
       np.testing.assert_allclose(result, force, rtol=1e-5, atol=2)
 
       # check for zeros after first condim elements
-      condim = dx.contact.dim[j]
+      condim = dx._impl.contact.dim[j]
       if condim < 6:
         np.testing.assert_allclose(force[condim:], 0, rtol=1e-5, atol=1e-5)
 
@@ -412,8 +412,8 @@ class SupportTest(parameterized.TestCase):
           ),
       )(mx, dx, j, True)
       # back to contact frame
-      force = force.at[:3].set(dx.contact.frame[j] @ force[:3])
+      force = force.at[:3].set(dx._impl.contact.frame[j] @ force[:3])
-      force = force.at[3:].set(dx.contact.frame[j] @ force[3:])
+      force = force.at[3:].set(dx._impl.contact.frame[j] @ force[3:])
       np.testing.assert_allclose(result, force, rtol=1e-5, atol=2)
 
   def test_wrap_inside(self):

mjx/mujoco/mjx/_src/test_util.py

@@ -106,20 +106,20 @@ def benchmark(
 
 
 def efc_order(m: mujoco.MjModel, d: mujoco.MjData, dx: Data) -> np.ndarray:
-  """Returns a sort order such that dx.efc_*[order][:d.nefc] == d.efc_*."""
+  """Returns a sort order such that dx.efc_*[order][:d._impl.nefc] == d.efc_*."""  # pytype: disable=attribute-error
   # reorder efc rows to skip inactive constraints and match contact order
-  efl = dx.ne + dx.nf + dx.nl
+  efl = dx._impl.ne + dx._impl.nf + dx._impl.nl  # pytype: disable=attribute-error
   order = np.arange(efl)
-  order[(dx.efc_J[:efl] == 0).all(axis=1)] = 2**16  # move empty rows to end
+  order[(dx._impl.efc_J[:efl] == 0).all(axis=1)] = 2**16  # move empty rows to end  # pytype: disable=attribute-error
-  for i in range(dx.ncon):
+  for i in range(dx._impl.ncon):  # pytype: disable=attribute-error
-    num_rows = dx.contact.dim[i]
+    num_rows = dx._impl.contact.dim[i]  # pytype: disable=attribute-error
-    if dx.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
+    if dx._impl.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:  # pytype: disable=attribute-error
-      num_rows = (dx.contact.dim[i] - 1) * 2
+      num_rows = (dx._impl.contact.dim[i] - 1) * 2  # pytype: disable=attribute-error
-    if dx.contact.dist[i] > 0:  # move empty contacts to end
+    if dx._impl.contact.dist[i] > 0:  # move empty contacts to end  # pytype: disable=attribute-error
       order = np.append(order, np.repeat(2**16, num_rows))
       continue
-    contact_match = (d.contact.geom == dx.contact.geom[i]).all(axis=-1)
+    contact_match = (d.contact.geom == dx._impl.contact.geom[i]).all(axis=-1)  # pytype: disable=attribute-error
-    contact_match &= (d.contact.pos == dx.contact.pos[i]).all(axis=-1)
+    contact_match &= (d.contact.pos == dx._impl.contact.pos[i]).all(axis=-1)  # pytype: disable=attribute-error
     assert contact_match.any(), f'contact {i} not found'
     contact_id = np.nonzero(contact_match)[0][0]
     order = np.append(order, np.repeat(efl + contact_id, num_rows))