From a6a9a850076e567bb2191c483546d021cdf97501 Mon Sep 17 00:00:00 2001 From: Google DeepMind Date: Mon, 12 Jan 2026 05:53:23 -0800 Subject: [PATCH] Implement get_data_into_cpp and make_data_cpp. PiperOrigin-RevId: 855192298 Change-Id: I9f9fd07094c40b9df1b645c2c893bb2605941e32 --- mjx/mujoco/mjx/_src/io.py | 103 +++++++++++++++++++++++++++++++-- mjx/mujoco/mjx/_src/io_test.py | 23 ++++---- 2 files changed, 110 insertions(+), 16 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index abf5f952..9e3c1d5e 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index cd76b40d..927f9d9c 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/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: