update mjx for naccdmax

PiperOrigin-RevId: 958479874
Change-Id: If0f48f2f372955f2948cb286966f5fce52c3f0ba

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.'

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