chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

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

Browse Source 
PiperOrigin-RevId: 874622402
Change-Id: I38a00633dea77d940e4b0422f93e9456ad84daf9
This commit is contained in:
Google DeepMind
2026-02-24 08:16:39 -08:00
committed by Copybara-Service
parent 1d6ff2cece
commit 94acdb33bd
3 changed files with 80 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+67 -22
View File
@@ -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
mjx/mujoco/mjx/_src/io_test.py
+13 -14
View File
@@ -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)
mjx/mujoco/mjx/_src/types.py
-2
View File
@@ -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):