jaxsim 0.4.3.dev305__tar.gz → 0.4.3.dev327__tar.gz

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/.github/workflows/ci_cd.yml

@@ -8,8 +8,7 @@ on:
     types:
       - published
   schedule:
-  # * is a special character in YAML so you have to quote this string
-  # Execute a "nightly" build at 2 AM UTC
+  # Execute a nightly build at 2am UTC.
   - cron:  '0 2 * * *'
 
 
@@ -138,6 +137,8 @@ jobs:
            steps.changes.outputs.all == 'true')
         run: pytest
         env:
+          # https://github.com/pytest-dev/pytest/issues/7443#issuecomment-656642591
+          PY_COLORS: "1"
           JAX_PLATFORM_NAME: cpu
 
   publish:

jaxsim-0.4.3.dev305/.github/workflows/update_pixi_lockfile.yml → jaxsim-0.4.3.dev327/.github/workflows/pixi.yml

@@ -1,4 +1,4 @@
-name: Update pixi lockfile
+name: Pixi
 
 permissions:
   contents: write
@@ -7,12 +7,16 @@ permissions:
 on:
   workflow_dispatch:
   schedule:
-    - cron: 0 5 1 * *
+    # Execute at 5am UTC on the first day of the month.
+    - cron: '0 5 1 * *'
 
 jobs:
+
   pixi-update:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
+
     steps:
+
       - uses: actions/checkout@v4
         with:
           lfs: true
@@ -30,9 +34,14 @@ jobs:
           set -o pipefail
           pixi update --json | pixi exec pixi-diff-to-markdown --explicit-column > diff.md
 
+      - name: Test project against updated pixi
+        run: pixi run --environment tasks-cpu tests
+        env:
+          PY_COLORS: "1"
+          JAX_PLATFORM_NAME: cpu
+
       - name: Commit and push changes
-        run: |
-          echo "BRANCH_NAME=update-pixi-$(date +'%Y%m%d%H%M%S')" >> $GITHUB_ENV
+        run: echo "BRANCH_NAME=update-pixi-$(date +'%Y%m%d%H%M%S')" >> $GITHUB_ENV
 
       - name: Create pull request
         uses: peter-evans/create-pull-request@v7

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev305
+Version: 0.4.3.dev327
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>
@@ -51,6 +51,7 @@ Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Topic :: Games/Entertainment :: Simulation
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/pyproject.toml

@@ -37,6 +37,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Programming Language :: Python :: Implementation :: CPython",
     "Topic :: Games/Entertainment :: Simulation",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
@@ -132,8 +133,10 @@ testpaths = [
 [tool.ruff]
 exclude = [
     ".git",
+    ".pixi",
     ".pytest_cache",
     ".ruff_cache",
+    ".idea",
     ".vscode",
     ".devcontainer",
     "__pycache__",
@@ -179,52 +182,87 @@ ignore = [
 channels = ["conda-forge"]
 platforms = ["linux-64", "osx-arm64", "osx-64"]
 
+[tool.pixi.environments]
+# We resolve only two groups: cpu and gpu.
+# Then, multiple environments can be created from these groups.
+default = { solve-group = "cpugroup" }
+gpu = { features = ["gpu"], solve-group = "gpugroup" }
+tasks-cpu = { features = ["test", "examples"], solve-group = "cpugroup" }
+tasks-gpu = { features = ["test", "examples", "gpu"], solve-group = "gpugroup" }
+
+# ---------------
+# feature.default
+# ---------------
+
+# Dependencies from conda-forge.
 [tool.pixi.dependencies]
+#
+# Matching `project.dependencies`.
+#
 coloredlogs = "*"
 jax = "*"
-jax-dataclasses = "*"
-jaxopt = "*"
 jaxlib = "*"
 jaxlie = "*"
-lxml = "*"
-mediapy = "*"
-mujoco = "*"
-notebook = "*"
+jax-dataclasses = "*"
 pptree = "*"
+optax = "*"
 qpax = "*"
 rod = "*"
-sdformat14 = "*"
 typing_extensions = "*"
+#
+# Optional dependencies.
+#
+lxml = "*"
+mediapy = "*"
+mujoco = "*"
+scipy = "*"
+#
+# Additional dependencies.
+#
+pip = "*"
+
+# Dependendencies from PyPI.
+[tool.pixi.pypi-dependencies]
+jaxsim = { path = "./", editable = true }
+
+# ------------
+# feature.test
+# ------------
 
 [tool.pixi.feature.test.tasks]
-examples = { cmd = "jupyter notebook ./examples" }
 pipcheck = "pip check"
-test = { cmd = "pytest", depends_on = ["pipcheck"] }
+tests = { cmd = "pytest", depends_on = ["pipcheck"] }
 
 [tool.pixi.feature.test.dependencies]
 black = "24.*"
 idyntree = "*"
 isort = "*"
-pip = "*"
 pre-commit = "*"
 pytest = "*"
 pytest-icdiff = "*"
 robot_descriptions = "*"
 
+# ----------------
+# feature.examples
+# ----------------
+
+[tool.pixi.feature.examples.tasks]
+examples = { cmd = "jupyter notebook ./examples" }
+
+[tool.pixi.feature.examples.dependencies]
+notebook = "*"
+robot_descriptions = "*"
+
+# -----------
+# feature.gpu
+# -----------
+
 [tool.pixi.feature.gpu]
 platforms = ["linux-64"]
-system-requirements = { cuda = "12.1" }
+system-requirements = { cuda = "12" }
 
 [tool.pixi.feature.gpu.dependencies]
-cuda-cupti = "*"
 cuda-version = "12.*"
-jaxlib = {version = "*", build = "*cuda*"}
-
-[tool.pixi.pypi-dependencies]
-jaxsim = { path = "./", editable = true }
-
-[tool.pixi.environments]
-default = { solve-group = "cpugroup" }
-gpu = { features = ["gpu"], solve-group = "gpugroup" }
-test-cpu = { features = ["test"], solve-group = "cpugroup" }
-test-gpu = { features = ["test", "gpu"], solve-group = "gpugroup" }
+# Pinning a specific version awaiting the following to get addressed:
+# https://github.com/conda-forge/jaxlib-feedstock/issues/285
+jaxlib = { version = "<0.4.31", build = "*cuda*" }

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev305'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev305')
+__version__ = version = '0.4.3.dev327'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev327')

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/api/contact.py

@@ -138,7 +138,7 @@ def collidable_point_dynamics(
     **kwargs,
 ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
     r"""
-    Compute the 6D force applied to each collidable point.
+    Compute the 6D force applied to each enabled collidable point.
 
     Args:
         model: The model to consider.
@@ -151,7 +151,7 @@ def collidable_point_dynamics(
         kwargs: Additional keyword arguments to pass to the active contact model.
 
     Returns:
-        The 6D force applied to each collidable point and additional data based
+        The 6D force applied to each eneabled collidable point and additional data based
         on the contact model configured:
         - Soft: the material deformation rate.
         - Rigid: no additional data.
@@ -199,15 +199,19 @@ def collidable_point_dynamics(
     )
 
     # Compute the transforms of the implicit frames `C[L] = (W_p_C, [L])`
-    # associated to each collidable point.
+    # associated to the enabled collidable point.
     # In inertial-fixed representation, the computation of these transforms
     # is not necessary and the conversion below becomes a no-op.
+
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
     W_H_C = (
         js.contact.transforms(model=model, data=data)
         if data.velocity_representation is not VelRepr.Inertial
-        else jnp.zeros(
-            shape=(len(model.kin_dyn_parameters.contact_parameters.body), 4, 4)
-        )
+        else jnp.zeros(shape=(len(indices_of_enabled_collidable_points), 4, 4))
     )
 
     # Convert the 6D forces to the active representation.
@@ -246,6 +250,15 @@ def in_contact(
     if link_names is not None and set(link_names).difference(model.link_names()):
         raise ValueError("One or more link names are not part of the model")
 
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
     W_p_Ci = collidable_point_positions(model=model, data=data)
 
     terrain_height = jax.vmap(lambda x, y: model.terrain.height(x=x, y=y))(
@@ -262,7 +275,7 @@ def in_contact(
 
     links_in_contact = jax.vmap(
         lambda link_index: jnp.where(
-            jnp.array(model.kin_dyn_parameters.contact_parameters.body) == link_index,
+            parent_link_idx_of_enabled_collidable_points == link_index,
             below_terrain,
             jnp.zeros_like(below_terrain, dtype=bool),
         ).any()
@@ -426,14 +439,14 @@ def estimate_good_contact_parameters(
 @jax.jit
 def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
     r"""
-    Return the pose of the collidable points.
+    Return the pose of the enabled collidable points.
 
     Args:
         model: The model to consider.
         data: The data of the considered model.
 
     Returns:
-        The stacked SE(3) matrices of all collidable points.
+        The stacked SE(3) matrices of all enabled collidable points.
 
     Note:
         Each collidable point is implicitly associated with a frame
@@ -442,16 +455,27 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
         rigidly attached to.
     """
 
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
     # Get the transforms of the parent link of all collidable points.
     W_H_L = js.model.forward_kinematics(model=model, data=data)[
-        jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
+        parent_link_idx_of_enabled_collidable_points
+    ]
+
+    L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
+        indices_of_enabled_collidable_points
     ]
 
     # Build the link-to-point transform from the displacement between the link frame L
     # and the implicit contact frame C.
-    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(
-        model.kin_dyn_parameters.contact_parameters.point
-    )
+    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(L_p_Ci)
 
     # Compose the work-to-link and link-to-point transforms.
     return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
@@ -465,7 +489,7 @@ def jacobian(
     output_vel_repr: VelRepr | None = None,
 ) -> jtp.Array:
     r"""
-    Return the free-floating Jacobian of the collidable points.
+    Return the free-floating Jacobian of the enabled collidable points.
 
     Args:
         model: The model to consider.
@@ -475,7 +499,7 @@ def jacobian(
 
     Returns:
         The stacked :math:`6 \times (6+n)` free-floating jacobians of the frames associated to the
-        collidable points.
+        enabled collidable points.
 
     Note:
         Each collidable point is implicitly associated with a frame
@@ -488,6 +512,15 @@ def jacobian(
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
     # Compute the Jacobians of all links.
     W_J_WL = js.model.generalized_free_floating_jacobian(
         model=model, data=data, output_vel_repr=VelRepr.Inertial
@@ -496,9 +529,7 @@ def jacobian(
     # Compute the contact Jacobian.
     # In inertial-fixed output representation, the Jacobian of the parent link is also
     # the Jacobian of the frame C implicitly associated with the collidable point.
-    W_J_WC = W_J_WL[
-        jnp.array(model.kin_dyn_parameters.contact_parameters.body, dtype=int)
-    ]
+    W_J_WC = W_J_WL[parent_link_idx_of_enabled_collidable_points]
 
     # Adjust the output representation.
     match output_vel_repr:
@@ -550,7 +581,7 @@ def jacobian_derivative(
     output_vel_repr: VelRepr | None = None,
 ) -> jtp.Matrix:
     r"""
-    Compute the derivative of the free-floating jacobian of the contact points.
+    Compute the derivative of the free-floating jacobian of the enabled collidable points.
 
     Args:
         model: The model to consider.
@@ -559,7 +590,7 @@ def jacobian_derivative(
             The output velocity representation of the free-floating jacobian derivative.
 
     Returns:
-        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the contact points.
+        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the enabled collidable points.
 
     Note:
         The input representation of the free-floating jacobian derivative is the active
@@ -570,10 +601,18 @@ def jacobian_derivative(
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
     # Get the index of the parent link and the position of the collidable point.
-    parent_link_idxs = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
-    L_p_Ci = jnp.array(model.kin_dyn_parameters.contact_parameters.point)
-    contact_idxs = jnp.arange(L_p_Ci.shape[0])
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
+        indices_of_enabled_collidable_points
+    ]
 
     # Get the transforms of all the parent links.
     W_H_Li = js.model.forward_kinematics(model=model, data=data)
@@ -646,7 +685,7 @@ def jacobian_derivative(
             output_vel_repr=VelRepr.Inertial,
         )
 
-    # Get the Jacobian of the collidable points in the mixed representation.
+    # Get the Jacobian of the enabled collidable points in the mixed representation.
     with data.switch_velocity_representation(VelRepr.Mixed):
         CW_J_WC_BW = jacobian(
             model=model,
@@ -656,13 +695,11 @@ def jacobian_derivative(
 
     def compute_O_J̇_WC_I(
         L_p_C: jtp.Vector,
-        contact_idx: jtp.Int,
+        parent_link_idx: jtp.Int,
         CW_J_WC_BW: jtp.Matrix,
         W_H_L: jtp.Matrix,
     ) -> jtp.Matrix:
 
-        parent_link_idx = parent_link_idxs[contact_idx]
-
         match output_vel_repr:
             case VelRepr.Inertial:
                 O_X_W = W_X_W = Adjoint.from_transform(  # noqa: F841
@@ -703,7 +740,7 @@ def jacobian_derivative(
         return O_J̇_WC_I
 
     O_J̇_WC = jax.vmap(compute_O_J̇_WC_I, in_axes=(0, 0, 0, None))(
-        L_p_Ci, contact_idxs, CW_J_WC_BW, W_H_Li
+        L_p_Ci, parent_link_idx_of_enabled_collidable_points, CW_J_WC_BW, W_H_Li
     )
 
     return O_J̇_WC

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/api/kin_dyn_parameters.py

@@ -743,6 +743,9 @@ class ContactParameters(JaxsimDataclass):
         point:
             The translations between the link frame and the collidable point, expressed
             in the coordinates of the parent link frame.
+        enabled:
+            A tuple of booleans representing, for each collidable point, whether it is
+            enabled or not in contact models.
 
     Note:
         Contrarily to LinkParameters and JointParameters, this class is not meant

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/api/model.py

@@ -2287,7 +2287,14 @@ def step(
                 msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
             )
 
-            W_p_C = js.contact.collidable_point_positions(model, data_tf)
+            # Extract the indices corresponding to the enabled collidable points.
+            indices_of_enabled_collidable_points = (
+                model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+            )
+
+            W_p_C = js.contact.collidable_point_positions(model, data_tf)[
+                indices_of_enabled_collidable_points
+            ]
 
             # Compute the penetration depth of the collidable points.
             δ, *_ = jax.vmap(
@@ -2296,8 +2303,9 @@ def step(
             )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
 
             with data_tf.switch_velocity_representation(VelRepr.Mixed):
-
-                J_WC = js.contact.jacobian(model, data_tf)
+                J_WC = js.contact.jacobian(model, data_tf)[
+                    indices_of_enabled_collidable_points
+                ]
                 M = js.model.free_floating_mass_matrix(model, data_tf)
 
                 # Compute the impact velocity.

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/rbda/collidable_points.py

@@ -21,7 +21,7 @@ def collidable_points_pos_vel(
 ) -> tuple[jtp.Matrix, jtp.Matrix]:
     """
 
-    Compute the position and linear velocity of collidable points in the world frame.
+    Compute the position and linear velocity of the enabled collidable points in the world frame.
 
     Args:
         model: The model to consider.
@@ -35,10 +35,23 @@ def collidable_points_pos_vel(
         joint_velocities: The velocities of the joints.
 
     Returns:
-        A tuple containing the position and linear velocity of collidable points.
+        A tuple containing the position and linear velocity of the enabled collidable points.
     """
 
-    if len(model.kin_dyn_parameters.contact_parameters.body) == 0:
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
+        indices_of_enabled_collidable_points
+    ]
+
+    if len(indices_of_enabled_collidable_points) == 0:
         return jnp.array(0).astype(float), jnp.empty(0).astype(float)
 
     W_p_B, W_Q_B, s, W_v_WB, ṡ, _, _, _, _, _ = utils.process_inputs(
@@ -136,8 +149,8 @@ def collidable_points_pos_vel(
 
     # Process all the collidable points in parallel.
     W_p_Ci, CW_vl_WC = jax.vmap(process_point_kinematics)(
-        model.kin_dyn_parameters.contact_parameters.point,
-        jnp.array(model.kin_dyn_parameters.contact_parameters.body),
+        L_p_Ci,
+        parent_link_idx_of_enabled_collidable_points,
     )
 
     return W_p_Ci, CW_vl_WC

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/rbda/contacts/common.py

@@ -216,11 +216,21 @@ class ContactModel(JaxsimDataclass):
             the velocity representation of data.
         """
 
+        # Get the object storing the contact parameters of the model.
+        contact_parameters = model.kin_dyn_parameters.contact_parameters
+
+        # Extract the indices corresponding to the enabled collidable points.
+        indices_of_enabled_collidable_points = (
+            contact_parameters.indices_of_enabled_collidable_points
+        )
+
         # Convert the contact forces to a JAX array.
         f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
 
         # Get the pose of the enabled collidable points.
-        W_H_C = js.contact.transforms(model=model, data=data)
+        W_H_C = js.contact.transforms(model=model, data=data)[
+            indices_of_enabled_collidable_points
+        ]
 
         # Convert the contact forces to inertial-fixed representation.
         W_f_C = jax.vmap(
@@ -234,14 +244,6 @@ class ContactModel(JaxsimDataclass):
             )
         )(f_C, W_H_C)
 
-        # Get the object storing the contact parameters of the model.
-        contact_parameters = model.kin_dyn_parameters.contact_parameters
-
-        # Extract the indices corresponding to the enabled collidable points.
-        indices_of_enabled_collidable_points = (
-            contact_parameters.indices_of_enabled_collidable_points
-        )
-
         # Construct the vector defining the parent link index of each collidable point.
         # We use this vector to sum the 6D forces of all collidable points rigidly
         # attached to the same link.

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/rbda/contacts/relaxed_rigid.py

@@ -357,13 +357,15 @@ class RelaxedRigidContacts(common.ContactModel):
 
             Jl_WC = jnp.vstack(
                 jax.vmap(lambda J, height: J * (height < 0))(
-                    js.contact.jacobian(model=model, data=data)[:, :3, :], δ
+                    js.contact.jacobian(model=model, data=data)[:, :3, :],
+                    δ,
                 )
             )
 
             J̇_WC = jnp.vstack(
                 jax.vmap(lambda J̇, height: J̇ * (height < 0))(
-                    js.contact.jacobian_derivative(model=model, data=data)[:, :3], δ
+                    js.contact.jacobian_derivative(model=model, data=data)[:, :3],
+                    δ,
                 ),
             )
 
@@ -530,6 +532,15 @@ class RelaxedRigidContacts(common.ContactModel):
             )
         )
 
+        # Get the indices of the enabled collidable points.
+        indices_of_enabled_collidable_points = (
+            model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+        )
+
+        parent_link_idx_of_enabled_collidable_points = jnp.array(
+            model.kin_dyn_parameters.contact_parameters.body, dtype=int
+        )[indices_of_enabled_collidable_points]
+
         # Compute the 6D inertia matrices of all links.
         M_L = js.model.link_spatial_inertia_matrices(model=model)
 
@@ -595,9 +606,7 @@ class RelaxedRigidContacts(common.ContactModel):
             f=jnp.concatenate,
             tree=(
                 *jax.vmap(compute_row)(
-                    link_idx=jnp.array(
-                        model.kin_dyn_parameters.contact_parameters.body
-                    ),
+                    link_idx=parent_link_idx_of_enabled_collidable_points,
                     penetration=penetration,
                     velocity=velocity,
                 ),

{jaxsim-0.4.3.dev305 → jaxsim-0.4.3.dev327}/src/jaxsim/rbda/contacts/rigid.py

@@ -285,6 +285,13 @@ class RigidContacts(ContactModel):
         # Import qpax privately just in this method.
         import qpax
 
+        # Get the indices of the enabled collidable points.
+        indices_of_enabled_collidable_points = (
+            model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+        )
+
+        n_collidable_points = len(indices_of_enabled_collidable_points)
+
         link_forces = jnp.atleast_2d(
             jnp.array(link_forces, dtype=float).squeeze()
             if link_forces is not None
@@ -299,7 +306,6 @@ class RigidContacts(ContactModel):
 
         # Compute kin-dyn quantities used in the contact model.
         with data.switch_velocity_representation(VelRepr.Mixed):
-
             BW_ν = data.generalized_velocity()
 
             M = js.model.free_floating_mass_matrix(model=model, data=data)
@@ -310,14 +316,11 @@ class RigidContacts(ContactModel):
             W_H_C = js.contact.transforms(model=model, data=data)
 
         # Compute the position and linear velocities (mixed representation) of
-        # all collidable points belonging to the robot.
+        # all enabled collidable points belonging to the robot.
         position, velocity = js.contact.collidable_point_kinematics(
             model=model, data=data
         )
 
-        # Get the number of collidable points.
-        n_collidable_points = len(model.kin_dyn_parameters.contact_parameters.body)
-
         # Compute the penetration depth and velocity of the collidable points.
         # Note that this function considers the penetration in the normal direction.
         δ, δ_dot, n̂ = jax.vmap(common.compute_penetration_data, in_axes=(0, 0, None))(
@@ -460,7 +463,7 @@ class RigidContacts(ContactModel):
         return G
 
     @staticmethod
-    def _compute_ineq_bounds(n_collidable_points: jtp.FloatLike) -> jtp.Vector:
+    def _compute_ineq_bounds(n_collidable_points: int) -> jtp.Vector:
 
         n_constraints = 6 * n_collidable_points
         return jnp.zeros(shape=(n_constraints,))