Fix Impl.CPP memory leaks due to mjModel/mjData not clearing from pytree.

PiperOrigin-RevId: 874622402
Change-Id: I38a00633dea77d940e4b0422f93e9456ad84daf9

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 keepalive_refs is None:
-
+    raise ValueError('keepalive_refs must be provided for CPP implementation.')
-  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
@@ -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
 

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)
+    lo = vmjx_data._impl.pointer_lo
-    # check that the data pointers in fact point to different datas
+    hi = vmjx_data._impl.pointer_hi
-    self.assertNotEqual(
+    addr0 = int(lo[0]) | (int(hi[0]) << 32)
-        vmjx_data._impl._data[0]._address,
+    addr1 = int(lo[1]) | (int(hi[1]) << 32)
-        vmjx_data._impl._data[1]._address,
+    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)
+    lo = vmjx_data._impl.pointer_lo
-    # check that the data pointers in fact point to different datas
+    hi = vmjx_data._impl.pointer_hi
-    self.assertNotEqual(
+    addr0 = int(lo[0]) | (int(hi[0]) << 32)
-        vmjx_data._impl._data[0]._address,
+    addr1 = int(lo[1]) | (int(hi[1]) << 32)
-        vmjx_data._impl._data[1]._address,
+    self.assertNotEqual(addr0, addr1)
-    )
 
   def test_get_set_state(self):
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)

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