Enable using sensor and SDF plugins in MJX:C.

PiperOrigin-RevId: 765089354
Change-Id: I3cacc2980ac8aeff68cfb927beb5b1dbb13261d4

mjx/mujoco/mjx/_src/collision_driver.py

@@ -363,8 +363,12 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
 
   condim_counts = {}
   for k, v in group_counts.items():
-    func = _COLLISION_FUNC[k.types]
+    if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF:
-    num_contacts = condim_counts.get(k.condim, 0) + func.ncon * v  # pytype: disable=attribute-error
+      ncon = m.opt.sdf_initpoints
+    else:
+      func = _COLLISION_FUNC[k.types]
+      ncon = func.ncon  # pytype: disable=attribute-error
+    num_contacts = condim_counts.get(k.condim, 0) + ncon * v
     if max_contact_points > -1:
       num_contacts = min(max_contact_points, num_contacts)
     condim_counts[k.condim] = num_contacts

mjx/mujoco/mjx/_src/io.py

@@ -626,6 +626,7 @@ def _make_data_c(
       'qLDiagInv': (m.nv, float_),
       'ten_velocity': (m.ntendon, float_),
       'actuator_velocity': (m.nu, float_),
+      'plugin_data': (m.nplugin, np.uint64),
       'B_rownnz': (m.nbody, np.int32),
       'B_rowadr': (m.nbody, np.int32),
       'B_colind': (m.nB, np.int32),

mjx/mujoco/mjx/_src/types.py

@@ -485,6 +485,7 @@ class Option(PyTreeNode):
   disableflags: DisableBit
   enableflags: int
   disableactuator: int
+  sdf_initpoints: int
 
 
 class OptionC(Option):
@@ -495,7 +496,6 @@ class OptionC(Option):
   ccd_tolerance: jax.Array
   noslip_iterations: int
   ccd_iterations: int
-  sdf_initpoints: int
   sdf_iterations: int
 
 
@@ -519,6 +519,7 @@ class ModelC(PyTreeNode):
   nflexshelldata: jax.Array
   nflexevpair: jax.Array
   nflextexcoord: jax.Array
+  nplugin: jax.Array
   ntree: jax.Array
   narena: jax.Array
   body_bvhadr: jax.Array
@@ -526,6 +527,7 @@ class ModelC(PyTreeNode):
   bvh_child: jax.Array
   bvh_nodeid: jax.Array
   bvh_aabb: jax.Array
+  geom_plugin: jax.Array
   light_bodyid: jax.Array
   light_targetbodyid: jax.Array
   flex_contype: jax.Array
@@ -569,6 +571,8 @@ class ModelC(PyTreeNode):
   flex_bvhadr: jax.Array
   flex_bvhnum: jax.Array
   actuator_plugin: jax.Array
+  sensor_plugin: jax.Array
+  plugin: jax.Array
 
 
 class ModelJAX(PyTreeNode):
@@ -980,6 +984,7 @@ class DataC(PyTreeNode):
   ten_velocity: jax.Array
   actuator_velocity: jax.Array
   cdof_dot: jax.Array
+  plugin_data: jax.Array
   qH: jax.Array  # pylint:disable=invalid-name
   qHDiagInv: jax.Array  # pylint:disable=invalid-name
   B_rownnz: jax.Array  # pylint:disable=invalid-name