diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 670c93e1..d4e0d1fa 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -494,9 +494,12 @@ def _put_model_warp( return _strip_weak_type(model) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _put_model_cpp( m: mujoco.MjModel, device: Optional[jax.Device] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model.""" @@ -517,9 +520,11 @@ def _put_model_cpp( c_pointers_impl = types.ModelCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _model=m, ) + if keepalive_refs is not None: + keepalive_refs[addr] = m + model = types.Model( **{k: copy.copy(v) for k, v in fields.items()}, _impl=c_pointers_impl ) @@ -532,6 +537,7 @@ def put_model( device: Optional[jax.Device] = None, impl: Optional[Union[str, types.Impl]] = None, graph_mode: Optional[mjxw.types.GraphMode] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Model: """Puts mujoco.MjModel onto a device, resulting in mjx.Model. @@ -541,6 +547,9 @@ def put_model( impl: implementation to use graph_mode: CUDA graph capture mode (for Warp only). Use GraphMode enum from warp._src.jax_experimental.ffi. GraphMode.WARP is the default mode. + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Required for CPP + impl to keep the model alive. Returns: an mjx.Model placed on device @@ -560,7 +569,7 @@ def put_model( graph_mode = graph_mode or getattr(mjxw.types.GraphMode, 'WARP') return _put_model_warp(m, graph_mode, device) elif impl == types.Impl.CPP: - return _put_model_cpp(m, device) + return _put_model_cpp(m, device, keepalive_refs=keepalive_refs) else: raise ValueError(f'Unsupported implementation: {impl}') @@ -906,9 +915,12 @@ def _make_data_warp( return data +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _make_data_cpp( m: Union[types.Model, mujoco.MjModel], device: Optional[jax.Device] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Allocate and initialize Data for the CPP implementation.""" if isinstance(m, mujoco.MjModel): @@ -918,7 +930,13 @@ def _make_data_cpp( 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 + model_addr = int(m_impl.pointer_lo) | (int(m_impl.pointer_hi) << 32) + if keepalive_refs is None or model_addr not in keepalive_refs: + raise ValueError( + 'keepalive_refs must be provided and contain the model when calling' + ' _make_data_cpp with a types.Model.' + ) + mj_model = keepalive_refs[model_addr] # Create the raw MuJoCo data mj_data = mujoco.MjData(mj_model) @@ -933,9 +951,11 @@ def _make_data_cpp( c_pointers_impl = types.DataCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _data=[mj_data], ) + if keepalive_refs is not None: + keepalive_refs[addr] = [mj_data] + data = types.Data( _impl=c_pointers_impl, **fields, @@ -952,6 +972,7 @@ def make_data( nconmax: Optional[int] = None, naconmax: Optional[int] = None, njmax: Optional[int] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Allocate and initialize Data. @@ -969,6 +990,9 @@ def make_data( `naconmax` argument to set the upper bound for the number of contacts across all worlds, rather than the `nconmax` argument from MuJoCo Warp. njmax: maximum number of constraints to allocate for warp across all worlds + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Required for CPP + impl when passing a types.Model. Returns: an initialized mjx.Data placed on device @@ -997,7 +1021,7 @@ def make_data( elif impl == types.Impl.C: return _make_data_c(m, device) elif impl == types.Impl.CPP: - return _make_data_cpp(m, device) + return _make_data_cpp(m, device, keepalive_refs=keepalive_refs) elif impl == types.Impl.WARP: _check_warp_installed() naconmax = nconmax if naconmax is None else naconmax @@ -1300,11 +1324,14 @@ def _put_data_c( return _strip_weak_type(data) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _put_data_cpp( m: mujoco.MjModel, d: mujoco.MjData, device: Optional[jax.Device] = None, dummy_arg_for_batching: Optional[jax.Array] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" @@ -1317,6 +1344,8 @@ def _put_data_cpp( mujoco.mj_copyData(new_d, m, d) data_list.append(new_d) addr = new_d._address + if keepalive_refs is not None: + keepalive_refs[addr] = new_d # To ensure that we retain the full pointer even if jax.config.enable_x64 is # set to True, we store the pointer as two 32-bit values. In the FFI call, # we combine the two values into a single pointer value. @@ -1341,7 +1370,6 @@ def _put_data_cpp( c_pointers_impl = types.DataCPP( pointer_lo=pointer_lo, pointer_hi=pointer_hi, - _data=data_list, ) data = types.Data( @@ -1399,6 +1427,7 @@ def put_data( naconmax: Optional[int] = None, njmax: Optional[int] = None, dummy_arg_for_batching: Optional[jax.Array] = None, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> types.Data: """Puts mujoco.MjData onto a device, resulting in mjx.Data. @@ -1415,6 +1444,8 @@ def put_data( njmax: maximum number of constraints to allocate for warp dummy_arg_for_batching: dummy argument to use for batching in cpp implementation + keepalive_refs: optional dict to store references to underlying MuJoCo + objects, preventing them from being garbage collected. Returns: an mjx.Data placed on device @@ -1434,7 +1465,11 @@ def put_data( return _put_data_c(m, d, device) elif impl == types.Impl.CPP: return _put_data_cpp( - m, d, device, dummy_arg_for_batching=dummy_arg_for_batching + m, + d, + device, + dummy_arg_for_batching=dummy_arg_for_batching, + keepalive_refs=keepalive_refs, ) elif impl == types.Impl.WARP: _check_warp_installed() @@ -1692,15 +1727,18 @@ def _get_data_into( mujoco.mj_factorM(m, result_i) +# TODO(josechenf): Iterate on the keepalive implementation to make it easier to +# use before OSS. def _get_data_into_cpp( result: Union[mujoco.MjData, List[mujoco.MjData]], m: mujoco.MjModel, d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ): """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 + objects that are stored in keepalive_refs. This function simply copies the data from those underlying MjData objects to the result using mj_copyData. """ @@ -1712,13 +1750,8 @@ def _get_data_into_cpp( 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.' - ) + if keepalive_refs is None: + raise ValueError('keepalive_refs must be provided for CPP implementation.') # 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 @@ -1728,7 +1761,19 @@ def _get_data_into_cpp( d_i: types.Data = ( jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d ) - src_data = mj_data_list[i] + result_i = result[i] if batched else result + + if batched: + addr_i = int(d_impl.pointer_lo[i]) | (int(d_impl.pointer_hi[i]) << 32) + else: + addr_i = int(d_impl.pointer_lo) | (int(d_impl.pointer_hi) << 32) + + if addr_i not in keepalive_refs: + raise ValueError( + f'Address {addr_i} not found in keepalive_refs. ' + 'Ensure keepalive_refs from the original compile() is passed.' + ) + src_data = keepalive_refs[addr_i] needs_syncing = False for field in fields_to_check: @@ -1745,9 +1790,6 @@ def _get_data_into_cpp( src_data.mocap_quat[:] = d_i.mocap_quat mujoco.mj_kinematics(m, src_data) - 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) @@ -1755,6 +1797,7 @@ def get_data_into( result: Union[mujoco.MjData, List[mujoco.MjData]], m: mujoco.MjModel, d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ): """Gets mjx.Data from a device into an existing mujoco.MjData or list.""" is_batched = isinstance(result, list) @@ -1770,7 +1813,7 @@ def get_data_into( return _get_data_into(result, m, d) if d.impl == types.Impl.CPP: - return _get_data_into_cpp(result, m, d) + return _get_data_into_cpp(result, m, d, keepalive_refs=keepalive_refs) if d.impl == types.Impl.WARP: return _get_data_into_warp(result, m, d) @@ -1781,7 +1824,9 @@ def get_data_into( def get_data( - m: mujoco.MjModel, d: types.Data + m: mujoco.MjModel, + d: types.Data, + keepalive_refs: Optional[Dict[int, Any]] = None, ) -> Union[mujoco.MjData, List[mujoco.MjData]]: """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData].""" batched = len(d.qpos.shape) > 1 @@ -1792,7 +1837,7 @@ def get_data( else: result = mujoco.MjData(m) - get_data_into(result, m, d) + get_data_into(result, m, d, keepalive_refs=keepalive_refs) return result diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 1cd0afe9..9910e9b5 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -768,9 +768,10 @@ class DataIOTest(parameterized.TestCase): m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) d = mujoco.MjData(m) mujoco.mj_step(m, d, 2) - dx = mjx.put_data(m, d, impl=impl) + keepalive = {} if impl == 'cpp' else None + dx = mjx.put_data(m, d, impl=impl, keepalive_refs=keepalive) d_2 = mujoco.MjData(m) - mjx.get_data_into(d_2, m, dx) + mjx.get_data_into(d_2, m, dx, keepalive_refs=keepalive) # check a few fields np.testing.assert_allclose(d_2.qpos, d.qpos) @@ -931,12 +932,11 @@ class DataIOTest(parameterized.TestCase): ) self.assertEqual(vmjx_data.qpos.shape, (2, m.nq)) - self.assertEqual(len(vmjx_data._impl._data), 2) - # check that the data pointers in fact point to different datas - self.assertNotEqual( - vmjx_data._impl._data[0]._address, - vmjx_data._impl._data[1]._address, - ) + lo = vmjx_data._impl.pointer_lo + hi = vmjx_data._impl.pointer_hi + addr0 = int(lo[0]) | (int(hi[0]) << 32) + addr1 = int(lo[1]) | (int(hi[1]) << 32) + self.assertNotEqual(addr0, addr1) # Test cases for `_resolve_impl_and_device` where the device is @@ -1265,12 +1265,11 @@ class StateIOTest(parameterized.TestCase): ) self.assertEqual(vmjx_data.qpos.shape, (2, m.nq)) - self.assertEqual(len(vmjx_data._impl._data), 2) - # check that the data pointers in fact point to different datas - self.assertNotEqual( - vmjx_data._impl._data[0]._address, - vmjx_data._impl._data[1]._address, - ) + lo = vmjx_data._impl.pointer_lo + hi = vmjx_data._impl.pointer_hi + addr0 = int(lo[0]) | (int(hi[0]) << 32) + addr1 = int(lo[1]) | (int(hi[1]) << 32) + self.assertNotEqual(addr0, addr1) def test_get_set_state(self): m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index fcdc41d6..9f25a845 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -538,7 +538,6 @@ class ModelCPP(PyTreeNode): # we combine the two values into a single pointer value. pointer_lo: jax.Array pointer_hi: jax.Array - _model: mujoco.MjModel class DataCPP(PyTreeNode): @@ -548,7 +547,6 @@ class DataCPP(PyTreeNode): # we combine the two values into a single pointer value. pointer_lo: jax.Array pointer_hi: jax.Array - _data: list[Any] = dataclasses.field(default_factory=list, repr=False) class ModelC(PyTreeNode):