Implement get_data_into_cpp and make_data_cpp.

PiperOrigin-RevId: 855192298
Change-Id: I9f9fd07094c40b9df1b645c2c893bb2605941e32

mjx/mujoco/mjx/_src/io.py

@@ -912,6 +912,44 @@ def _make_data_warp(
   return data
 
 
+def _make_data_cpp(
+    m: Union[types.Model, mujoco.MjModel],
+    device: Optional[jax.Device] = None,
+) -> types.Data:
+  """Allocate and initialize Data for the CPP implementation."""
+  if isinstance(m, mujoco.MjModel):
+    mj_model = m
+  else:
+    # Get the underlying MjModel from the types.Model
+    m_impl = m._impl  # pylint: disable=protected-access
+    if not isinstance(m_impl, types.ModelCPP):
+      raise ValueError(f'Expected ModelCPP impl, got {type(m_impl)}')
+    mj_model = m_impl._model  # pylint: disable=protected-access
+
+  # Create the raw MuJoCo data
+  mj_data = mujoco.MjData(mj_model)
+
+  # Get the pointer address
+  addr = mj_data._address  # pytype: disable=attribute-error
+  pointer_lo = jp.array(addr & 0xFFFFFFFF, dtype=jp.uint32)
+  pointer_hi = jp.array(addr >> 32, dtype=jp.uint32)
+
+  fields = _put_data_public_fields(mj_data)
+
+  c_pointers_impl = types.DataCPP(
+      pointer_lo=pointer_lo,
+      pointer_hi=pointer_hi,
+      _data=[mj_data],
+  )
+
+  data = types.Data(
+      _impl=c_pointers_impl,
+      **fields,
+  )
+  data = jax.device_put(data, device=device)
+  return _strip_weak_type(data)
+
+
 def make_data(
     m: Union[types.Model, mujoco.MjModel],
     device: Optional[jax.Device] = None,
@@ -965,6 +1003,8 @@ def make_data(
     return _make_data_jax(m, device)
   elif impl == types.Impl.C:
     return _make_data_c(m, device)
+  elif impl == types.Impl.CPP:
+    return _make_data_cpp(m, device)
   elif impl == types.Impl.WARP:
     naconmax = nconmax if naconmax is None else naconmax
     return _make_data_warp(m, device, naconmax, njmax)
@@ -1265,12 +1305,8 @@ def _put_data_cpp(
   def _copy_and_get_addr(unused_jax_array):
     # We use the input to the callback as a dummy dependency to ensure
     # io_callback runs for each element in the batch.
-    try:
-      new_d = mujoco.MjData(m)
-    except mujoco.FatalError as e:
-      raise ValueError('Failed to create new MjData') from e
-    mujoco.mj_copyState(m, d, new_d, int(mujoco.mjtState.mjSTATE_FULLPHYSICS))
-    mujoco.mj_forward(m, new_d)
+    new_d = mujoco.MjData(m)
+    mujoco.mj_copyData(new_d, m, d)
     data_list.append(new_d)
     addr = new_d._address
     # To ensure that we retain the full pointer even if jax.config.enable_x64 is
@@ -1586,6 +1622,58 @@ def _get_data_into(
     mujoco.mj_factorM(m, result_i)
 
 
+def _get_data_into_cpp(
+    result: Union[mujoco.MjData, List[mujoco.MjData]],
+    m: mujoco.MjModel,
+    d: types.Data,
+):
+  """Gets mjx.Data from CPP impl into an existing mujoco.MjData or list.
+
+  For the CPP implementation, the mjx.Data wraps underlying mujoco.MjData
+  objects that are stored in DataCPP._data. This function simply copies the
+  data from those underlying MjData objects to the result using mj_copyData.
+  """
+
+  batched = isinstance(result, list)
+  d = jax.device_get(d)
+  batch_size = d.qpos.shape[0] if batched else 1
+
+  d_impl = d._impl  # pylint: disable=protected-access
+  if not isinstance(d_impl, types.DataCPP):
+    raise ValueError(f'Expected DataCPP impl, got {type(d_impl)}')
+
+  mj_data_list = d_impl._data  # pylint: disable=protected-access
+
+  if batch_size > len(mj_data_list):
+    raise ValueError(
+        f'Batch size {batch_size} exceeds number of underlying MjData objects '
+        f'({len(mj_data_list)}). Cannot copy data.'
+    )
+
+  # Verify that the underlying MjData state matches the mjx.Data state
+  # Ideally we'd use mj_getState and get_state here but that requires an
+  # mjx.Model which we don't have access to in this function.
+  fields_to_check = ['qpos', 'qvel', 'act', 'mocap_pos', 'mocap_quat']
+  for i in range(batch_size):
+    d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
+    src_data = mj_data_list[i]
+
+    for field in fields_to_check:
+      mj_value = getattr(src_data, field)
+      mjx_value = np.asarray(getattr(d_i, field))
+      if not np.allclose(mj_value, mjx_value):
+        raise ValueError(
+            f'State mismatch at batch index {i}, field {field}: underlying '
+            'MjData does not match mjx.Data. The mjx.Data may have been '
+            'modified without updating the underlying MjData.'
+        )
+
+  for i in range(batch_size):
+    result_i = result[i] if batched else result
+    src_data = mj_data_list[i]
+    mujoco.mj_copyData(result_i, m, src_data)
+
+
 def get_data_into(
     result: Union[mujoco.MjData, List[mujoco.MjData]],
     m: mujoco.MjModel,
@@ -1604,6 +1692,9 @@ def get_data_into(
     # TODO(stunya): Split out _get_data_into once codepaths diverge enough.
     return _get_data_into(result, m, d)
 
+  if d.impl == types.Impl.CPP:
+    return _get_data_into_cpp(result, m, d)
+
   if d.impl == types.Impl.WARP:
     return _get_data_into_warp(result, m, d)
 

mjx/mujoco/mjx/_src/io_test.py

@@ -371,7 +371,7 @@ class DataIOTest(parameterized.TestCase):
       self.tempdir = tempfile.TemporaryDirectory()
       wp.config.kernel_cache_dir = self.tempdir.name
 
-  @parameterized.parameters('jax', 'c')
+  @parameterized.parameters('jax', 'c', 'cpp')
   def test_make_data(self, impl: str):
     """Test that make_data returns the correct shapes."""
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
@@ -384,7 +384,7 @@ class DataIOTest(parameterized.TestCase):
     nv = 19
     nefc = 185
 
-    self.assertEqual(d._impl.nefc, nefc)
+    # Check public Data fields that exist for all impls
     self.assertEqual(d.qpos.shape, (nq,))
     self.assertEqual(d.qvel.shape, (nv,))
     self.assertEqual(d.act.shape, (0,))
@@ -406,17 +406,20 @@ class DataIOTest(parameterized.TestCase):
     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.actuator_length.shape, (1,))
+
+    if impl == 'cpp':
+      self.assertTrue(hasattr(d._impl, 'pointer_lo'))
+      self.assertTrue(hasattr(d._impl, 'pointer_hi'))
+      return  # cpp does not populate other _impl fields
+
+    self.assertEqual(d._impl.nefc, nefc)
     self.assertEqual(d._impl.cinert.shape, (nbody, 10))
     self.assertEqual(d._impl.crb.shape, (nbody, 10))
-    self.assertEqual(d.actuator_length.shape, (1,))
     if impl == 'jax':
       self.assertEqual(d._impl.actuator_moment.shape, (1, nv))
     elif impl == 'c':
       self.assertEqual(d._impl.actuator_moment.shape, (m.nJmom,))
-    elif impl == 'cpp':
-      self.assertTrue(hasattr(d._impl, 'pointer_lo'))
-      self.assertTrue(hasattr(d._impl, 'pointer_hi'))
-      return  # cpp does not populate other fields in _impl
     self.assertEqual(d._impl.contact.dist.shape, (ncon,))
     self.assertEqual(d._impl.contact.pos.shape, (ncon, 3))
     self.assertEqual(d._impl.contact.frame.shape, (ncon, 3, 3))
@@ -472,7 +475,7 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d._impl.contact__dist.shape[0], 9)
     self.assertEqual(d._impl.efc__pos.shape[0], 23)
 
-  @parameterized.parameters('jax', 'c')
+  @parameterized.parameters('jax', 'c', 'cpp')
   def test_put_data(self, impl: str):
     """Test that put_data puts the correct data for dense and sparse."""
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
@@ -699,7 +702,7 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(ds[0].ncon, 1)
     self.assertEqual(ds[1].ncon, 0)
 
-  @parameterized.parameters('jax', 'c')
+  @parameterized.parameters('jax', 'c', 'cpp')
   def test_get_data_into(self, impl):
     """Test that get_data_into correctly populates an MjData."""
 
@@ -809,7 +812,7 @@ class DataIOTest(parameterized.TestCase):
 
     _ = jax.tree.map_with_path(check_ndim, dx)
 
-  @parameterized.parameters('jax', 'warp')
+  @parameterized.parameters('jax', 'warp', 'cpp')
   def test_data_slice(self, impl):
     """Tests that slice on Data works as expected."""
     if impl == 'warp' and not mjxw.WARP_INSTALLED: