diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 04a4a3d6..2f45bca8 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/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] - num_contacts = condim_counts.get(k.condim, 0) + func.ncon * v # pytype: disable=attribute-error + if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF: + 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 diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index c5b94513..f8aaf565 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/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), diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 3beb6b54..a60a1dd6 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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