From 382af63dc3ac15dc9ea76562d2be0b90ae025e03 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Mon, 3 Aug 2026 11:19:19 -0700 Subject: [PATCH] update mjx for naccdmax PiperOrigin-RevId: 958479874 Change-Id: If0f48f2f372955f2948cb286966f5fce52c3f0ba --- mjx/mujoco/mjx/_src/io.py | 24 ++++++++++++++++++++++-- mjx/mujoco/mjx/_src/io_test.py | 8 +++++--- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 09165836..480db62c 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -1167,6 +1167,7 @@ def _put_data_warp( d: mujoco.MjData, device: Optional[jax.Device] = None, naconmax: Optional[int] = None, + naccdmax: Optional[int] = None, njmax: Optional[int] = None, nvmax: Optional[int] = None, ) -> types.Data: @@ -1174,7 +1175,13 @@ def _put_data_warp( with wp.ScopedDevice('cpu'): # pylint: disable=undefined-variable dw = mjwp.put_data( - m, d, nworld=1, naconmax=naconmax, njmax=njmax, nvmax=nvmax + m, + d, + nworld=1, + naconmax=naconmax, + naccdmax=naccdmax, + njmax=njmax, + nvmax=nvmax, ) # pylint: disable=undefined-variable fields = _put_data_public_fields(d) @@ -1209,6 +1216,7 @@ def put_data( device: Optional[jax.Device] = None, impl: Optional[Union[str, types.Impl]] = None, naconmax: Optional[int] = None, + naccdmax: Optional[int] = None, njmax: Optional[int] = None, nvmax: Optional[int] = None, dummy_arg_for_batching: Optional[jax.Array] = None, @@ -1225,6 +1233,10 @@ def put_data( Since the number of worlds is **not** pre-defined in JAX, we use the `naconmax` argument to set the upper bound for the number of contacts across all worlds. + naccdmax: maximum number of contacts for GJK collision detection across all + worlds. Since the number of worlds is **not** pre-defined in JAX, we use the + `naccdmax` argument to set the upper bound for the number of contacts + across all worlds, rather than the `nccdmax` argument from MuJoCo Warp. njmax: maximum number of constraints to allocate for warp nvmax: capacity for compacted active DOFs per world dummy_arg_for_batching: dummy argument to use for batching in cpp @@ -1249,7 +1261,15 @@ def put_data( ) elif impl == types.Impl.WARP: _check_warp_installed() - return _put_data_warp(m, d, device, naconmax, njmax, nvmax) + return _put_data_warp( + m, + d, + device=device, + naconmax=naconmax, + naccdmax=naccdmax, + njmax=njmax, + nvmax=nvmax, + ) raise NotImplementedError( f'put_data for implementation "{impl}" not implemented yet.' diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 3fc2905e..eabc4fd9 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -532,19 +532,21 @@ class DataIOTest(parameterized.TestCase): if not mjx_io.has_cuda_gpu_device(): self.skipTest('No CUDA GPU device.') m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) - d = mjx.make_data(m, impl='warp', naconmax=9, njmax=23, nvmax=8) + d = mjx.make_data(m, impl='warp', naconmax=9, naccdmax=7, njmax=23, nvmax=8) self.assertEqual(d._impl.contact__dist.shape[0], 9) + self.assertEqual(d._impl.naccdmax, 7) self.assertEqual(d._impl.efc__pos.shape[0], 23) self.assertEqual(d._impl.nvmax, 8) - def test_put_data_warp_nvmax(self): + def test_put_data_warp(self): if not mjxw.WARP_INSTALLED: self.skipTest('Warp is not installed.') if not mjx_io.has_cuda_gpu_device(): self.skipTest('No CUDA GPU device.') m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS) d = mujoco.MjData(m) - dx = mjx.put_data(m, d, impl='warp', nvmax=8) + dx = mjx.put_data(m, d, impl='warp', naccdmax=7, nvmax=8) + self.assertEqual(dx._impl.naccdmax, 7) self.assertEqual(dx._impl.nvmax, 8) @parameterized.parameters('jax', 'cpp', 'warp')