jaxsim 0.2.1.dev128__tar.gz → 0.2.1.dev155__tar.gz

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/.github/workflows/ci_cd.yml

@@ -118,8 +118,7 @@ jobs:
               - *src
               - *tests
 
-      # https://gazebosim.org/docs/harmonic/install_ubuntu
-      - name: Install Gazebo Sim
+      - name: Install 'gz sdf' system command
         if: |
           contains(matrix.os, 'ubuntu') &&
           (github.event_name != 'pull_request' ||
@@ -130,7 +129,7 @@ jobs:
           sudo wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg
           echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null
           sudo apt-get update
-          sudo apt-get install gz-harmonic
+          sudo apt-get install --no-install-recommends libsdformat13 gz-tools2
 
       - name: Run the Python tests
         if: |

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/.gitignore

@@ -140,5 +140,10 @@ src/jaxsim/_version.py
 
 # ruff
 .ruff_cache/
+
 # pixi environments
 .pixi
+
+# data
+.mp4
+.png

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev128
+Version: 0.2.1.dev155
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it
@@ -68,7 +68,7 @@ Requires-Dist: mujoco>=3.0.0; extra == "all"
 
 JaxSim is a **differentiable physics engine** and **multibody dynamics library** designed for applications in control and robot learning, implemented with JAX.
 
-Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence. 
+Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence.
 
 ## Features
 
@@ -91,7 +91,7 @@ Its design facilitates research and accelerates prototyping in the intersection
 
 ### JaxSim as a multibody dynamics library
 
-- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians. 
+- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians.
 - Provides all the quantities included in the Euler-Poincarè formulation of the equations of motion.
 - Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].
 - Exposes all the necessary quantities to develop controllers in centroidal coordinates.
@@ -198,10 +198,10 @@ The main differences between MJX/Brax and JaxSim are as follows:
 
 - JaxSim supports out-of-the-box all SDF models with [Pose Frame Semantics][PFS].
 - JaxSim only supports collisions between points rigidly attached to bodies and a compliant ground surface.
-  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous  
+  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous
   function of the state $(\mathbf{q}, \boldsymbol{\nu})$, it doesn't require running any optimization algorithm
   when stepping the simulation forward.
-- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators. 
+- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators.
 
 [brax]: https://github.com/google/brax
 [mjx]: https://mujoco.readthedocs.io/en/3.0.0/mjx.html

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/README.md

@@ -2,7 +2,7 @@
 
 JaxSim is a **differentiable physics engine** and **multibody dynamics library** designed for applications in control and robot learning, implemented with JAX.
 
-Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence. 
+Its design facilitates research and accelerates prototyping in the intersection of robotics and artificial intelligence.
 
 ## Features
 
@@ -25,7 +25,7 @@ Its design facilitates research and accelerates prototyping in the intersection
 
 ### JaxSim as a multibody dynamics library
 
-- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians. 
+- Provides rigid body dynamics algorithms (RBDAs) like RNEA, ABA, CRBA, and Jacobians.
 - Provides all the quantities included in the Euler-Poincarè formulation of the equations of motion.
 - Supports body-fixed, inertial-fixed, and mixed [velocity representations][notation].
 - Exposes all the necessary quantities to develop controllers in centroidal coordinates.
@@ -132,10 +132,10 @@ The main differences between MJX/Brax and JaxSim are as follows:
 
 - JaxSim supports out-of-the-box all SDF models with [Pose Frame Semantics][PFS].
 - JaxSim only supports collisions between points rigidly attached to bodies and a compliant ground surface.
-  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous  
+  Our contact model requires careful tuning of its spring-damper parameters, but being an instantaneous
   function of the state $(\mathbf{q}, \boldsymbol{\nu})$, it doesn't require running any optimization algorithm
   when stepping the simulation forward.
-- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators. 
+- JaxSim mitigates the stiffness of the contact-aware system dynamics by providing variable-step integrators.
 
 [brax]: https://github.com/google/brax
 [mjx]: https://mujoco.readthedocs.io/en/3.0.0/mjx.html

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155/examples}/.gitattributes

@@ -1,3 +1,2 @@
 # GitHub syntax highlighting
 pixi.lock linguist-language=YAML
-

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/__init__.py

@@ -8,8 +8,9 @@ def _jnp_options() -> None:
 
     import jax
 
-    # Enable by default
-    if not ("JAX_ENABLE_X64" in os.environ and os.environ["JAX_ENABLE_X64"] == "0"):
+    # Enable by default 64bit precision in JAX.
+    if os.environ.get("JAX_ENABLE_X64", "1") != "0":
+
         logging.info("Enabling JAX to use 64bit precision")
         jax.config.update("jax_enable_x64", True)
 
@@ -27,6 +28,7 @@ def _np_options() -> None:
 
 
 def _is_editable() -> bool:
+
     import importlib.util
     import pathlib
     import site
@@ -45,11 +47,40 @@ def _is_editable() -> bool:
     return jaxsim_package_dir not in site.getsitepackages()
 
 
-# Initialize the logging verbosity
-if _is_editable():
-    logging.configure(level=logging.LoggingLevel.DEBUG)
-else:
-    logging.configure(level=logging.LoggingLevel.WARNING)
+def _get_default_logging_level(env_var: str) -> logging.LoggingLevel:
+    """
+    Get the default logging level.
+
+    Args:
+        env_var: The environment variable to check.
+
+    Returns:
+        The logging level to set.
+    """
+
+    import os
+
+    # Define the default logging level depending on the installation mode.
+    default_logging_level = (
+        logging.LoggingLevel.DEBUG
+        if _is_editable()  # noqa: F821
+        else logging.LoggingLevel.WARNING
+    )
+
+    # Allow to override the default logging level with an environment variable.
+    try:
+        return logging.LoggingLevel[
+            os.environ.get(env_var, default_logging_level.name).upper()
+        ]
+
+    except KeyError as exc:
+        msg = f"Invalid logging level defined in {env_var}='{os.environ[env_var]}'"
+        raise RuntimeError(msg) from exc
+
+
+# Configure the logger with the default logging level.
+logging.configure(level=_get_default_logging_level(env_var="JAXSIM_LOGGING_LEVEL"))
+
 
 # Configure JAX
 _jnp_options()
@@ -59,6 +90,7 @@ _np_options()
 
 del _jnp_options
 del _np_options
+del _get_default_logging_level
 del _is_editable
 
 from . import terrain  # isort:skip

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1.dev128'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev128')
+__version__ = version = '0.2.1.dev155'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev155')

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/api/contact.py

@@ -365,20 +365,20 @@ def jacobian(
 
             W_H_C = transforms(model=model, data=data)
 
-            def jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
+            def body_jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
                 C_X_W = jaxsim.math.Adjoint.from_transform(
                     transform=W_H_C, inverse=True
                 )
                 C_J_WC = C_X_W @ W_J_WC
                 return C_J_WC
 
-            O_J_WC = jax.vmap(jacobian)(W_H_C, W_J_WC)
+            O_J_WC = jax.vmap(body_jacobian)(W_H_C, W_J_WC)
 
         case VelRepr.Mixed:
 
             W_H_C = transforms(model=model, data=data)
 
-            def jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
+            def mixed_jacobian(W_H_C: jtp.Matrix, W_J_WC: jtp.Matrix) -> jtp.Matrix:
 
                 W_H_CW = W_H_C.at[0:3, 0:3].set(jnp.eye(3))
 
@@ -389,7 +389,7 @@ def jacobian(
                 CW_J_WC = CW_X_W @ W_J_WC
                 return CW_J_WC
 
-            O_J_WC = jax.vmap(jacobian)(W_H_C, W_J_WC)
+            O_J_WC = jax.vmap(mixed_jacobian)(W_H_C, W_J_WC)
 
         case _:
             raise ValueError(output_vel_repr)

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/api/kin_dyn_parameters.py

@@ -6,6 +6,7 @@ import jax.lax
 import jax.numpy as jnp
 import jax_dataclasses
 import jaxlie
+import numpy as np
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
@@ -220,7 +221,9 @@ class KynDynParameters(JaxsimDataclass):
             (
                 hash(self.number_of_links()),
                 hash(self.number_of_joints()),
-                hash(tuple(jnp.atleast_1d(self.parent_array).flatten().tolist())),
+                hash(tuple(np.atleast_1d(self.parent_array).flatten().tolist())),
+                hash(self._parent_array),
+                hash(self._support_body_array_bool),
             )
         )
 

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/api/link.py

@@ -241,8 +241,8 @@ def jacobian(
     )
 
     # Compute the actual doubly-left free-floating jacobian of the link.
-    κ = model.kin_dyn_parameters.support_body_array_bool[link_index]
-    B_J_WL_B = jnp.hstack([jnp.ones(5), κ]) * B_J_full_WX_B
+    κb = model.kin_dyn_parameters.support_body_array_bool[link_index]
+    B_J_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J_full_WX_B
 
     # Adjust the input representation such that `J_WL_I @ I_ν`.
     match data.velocity_representation:

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/api/model.py

@@ -17,7 +17,7 @@ import jaxsim.api as js
 import jaxsim.parsers.descriptions
 import jaxsim.typing as jtp
 from jaxsim.math import Cross
-from jaxsim.utils import HashlessObject, JaxsimDataclass, Mutability
+from jaxsim.utils import JaxsimDataclass, Mutability
 
 from .common import VelRepr
 
@@ -33,6 +33,7 @@ class JaxSimModel(JaxsimDataclass):
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
         default=jaxsim.terrain.FlatTerrain(), repr=False, compare=False, hash=False
     )
+
     kin_dyn_parameters: js.kin_dyn_parameters.KynDynParameters | None = (
         dataclasses.field(default=None, repr=False, compare=False, hash=False)
     )
@@ -41,13 +42,9 @@ class JaxSimModel(JaxsimDataclass):
         default=None, repr=False, compare=False, hash=False
     )
 
-    _description: Static[
-        HashlessObject[jaxsim.parsers.descriptions.ModelDescription | None]
-    ] = dataclasses.field(default=None, repr=False, compare=False, hash=False)
-
-    @property
-    def description(self) -> jaxsim.parsers.descriptions.ModelDescription:
-        return self._description.get()
+    description: Static[jaxsim.parsers.descriptions.ModelDescription | None] = (
+        dataclasses.field(default=None, repr=False, compare=False, hash=False)
+    )
 
     def __eq__(self, other: JaxSimModel) -> bool:
 
@@ -61,6 +58,7 @@ class JaxSimModel(JaxsimDataclass):
         return hash(
             (
                 hash(self.model_name),
+                hash(self.description),
                 hash(self.kin_dyn_parameters),
             )
         )
@@ -157,7 +155,7 @@ class JaxSimModel(JaxsimDataclass):
         # Build the model
         model = JaxSimModel(
             model_name=model_name,
-            _description=HashlessObject(obj=model_description),
+            description=model_description,
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
@@ -302,7 +300,7 @@ def reduce(
         locked_joint_positions:
             A dictionary containing the positions of the joints to be considered
             in the reduction process. The removed joints in the reduced model
-            will have their position locked to their value in this dictionary.
+            will have their position locked to their value of this dictionary.
             If a joint is not part of the dictionary, its position is set to zero.
     """
 
@@ -315,10 +313,9 @@ def reduce(
         new_joints = set(model.joint_names()) - set(locked_joint_positions)
         raise ValueError(f"Passed joints not existing in the model: {new_joints}")
 
-    # Copy the model description with a deep copy of the joints.
-    intermediate_description = dataclasses.replace(
-        model.description, joints=copy.deepcopy(model.description.joints)
-    )
+    # Operate on a deep copy of the model description in order to prevent problems
+    # when mutable attributes are updated.
+    intermediate_description = copy.deepcopy(model.description)
 
     # Update the initial position of the joints.
     # This is necessary to compute the correct pose of the link pairs connected
@@ -686,8 +683,6 @@ def forward_dynamics_aba(
         another representation C_v̇_WB expressed in a generic frame C.
         """
 
-        from jaxsim.math import Cross
-
         # In Mixed representation, we need to include a cross product in ℝ⁶.
         # In Inertial and Body representations, the cross product is always zero.
         C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
@@ -1483,12 +1478,7 @@ def link_bias_accelerations(
     # ================================================
 
     # Compute the base transform.
-    W_H_B = jaxlie.SE3.from_rotation_and_translation(
-        rotation=jaxlie.SO3.from_quaternion_xyzw(
-            xyzw=jaxsim.math.Quaternion.to_xyzw(wxyz=data.base_orientation())
-        ),
-        translation=data.base_position(),
-    ).as_matrix()
+    W_H_B = data.base_transform()
 
     def other_representation_to_inertial(
         C_v̇_WB: jtp.Vector, C_v_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WC: jtp.Vector
@@ -1529,9 +1519,12 @@ def link_bias_accelerations(
             W_H_C = W_H_BW
             with data.switch_velocity_representation(VelRepr.Mixed):
                 W_ṗ_B = data.base_velocity()[0:3]
-                W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)
+                BW_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)
+                W_X_BW = jaxsim.math.Adjoint.from_transform(transform=W_H_BW)
+                W_v_WC = W_v_W_BW = W_X_BW @ BW_v_W_BW
             with data.switch_velocity_representation(VelRepr.Mixed):
                 C_v_WB = BW_v_WB = data.base_velocity()
+
         case _:
             raise ValueError(data.velocity_representation)
 

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/api/ode.py

@@ -223,7 +223,9 @@ def system_velocity_dynamics(
 
 @jax.jit
 def system_position_dynamics(
-    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[jtp.Vector, jtp.Vector, jtp.Vector]:
     """
     Compute the dynamics of the system position.
@@ -231,6 +233,8 @@ def system_position_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
+        baumgarte_quaternion_regularization:
+            The Baumgarte regularization coefficient for adjusting the quaternion norm.
 
     Returns:
         A tuple containing the derivative of the base position, the derivative of the
@@ -250,6 +254,7 @@ def system_position_dynamics(
         quaternion=W_Q_B,
         omega=W_ω_WB,
         omega_in_body_fixed=False,
+        K=baumgarte_quaternion_regularization,
     ).squeeze()
 
     return W_ṗ_B, W_Q̇_B, ṡ
@@ -262,6 +267,7 @@ def system_dynamics(
     *,
     joint_forces: jtp.Vector | None = None,
     link_forces: jtp.Vector | None = None,
+    baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[ODEState, dict[str, Any]]:
     """
     Compute the dynamics of the system.
@@ -271,6 +277,9 @@ def system_dynamics(
         data: The data of the considered model.
         joint_forces: The joint forces to apply.
         link_forces: The 6D forces to apply to the links.
+        baumgarte_quaternion_regularization:
+            The Baumgarte regularization coefficient used to adjust the norm of the
+            quaternion (only used in integrators not operating on the SO(3) manifold).
 
     Returns:
         A tuple with an `ODEState` object storing in each of its attributes the
@@ -287,7 +296,11 @@ def system_dynamics(
     )
 
     # Extract the velocities.
-    W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(model=model, data=data)
+    W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
+        model=model,
+        data=data,
+        baumgarte_quaternion_regularization=baumgarte_quaternion_regularization,
+    )
 
     # Create an ODEState object populated with the derivative of each leaf.
     # Our integrators, operating on generic pytrees, will be able to handle it

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/math/joint_model.py

@@ -34,8 +34,8 @@ class JointModel:
         already in a vectorized form. In other words, it cannot be created using vmap.
     """
 
-    λ_H_pre: jax.Array
-    suc_H_i: jax.Array
+    λ_H_pre: jtp.Array
+    suc_H_i: jtp.Array
 
     joint_dofs: Static[tuple[int, ...]]
     joint_names: Static[tuple[str, ...]]

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/mujoco/loaders.py

@@ -188,10 +188,9 @@ class RodModelToMjcf:
         )
 
         # If considered joints are passed, make sure that they are all part of the model.
-        if considered_joints - set([j.name for j in rod_model.joints()]):
-            extra_joints = set(considered_joints) - set(
-                [j.name for j in rod_model.joints()]
-            )
+        if considered_joints - {j.name for j in rod_model.joints()}:
+            extra_joints = set(considered_joints) - {j.name for j in rod_model.joints()}
+
             msg = f"Couldn't find the following joints in the model: '{extra_joints}'"
             raise ValueError(msg)
 
@@ -352,7 +351,7 @@ class RodModelToMjcf:
         # Set alpha=0 to the color of all collision elements
         for geometry_element in mujoco_element.findall(".//geom[@rgba]"):
             if geometry_element.attrib.get("name") in collision_names:
-                r, g, b, a = geometry_element.attrib["rgba"].split(" ")
+                r, g, b, _ = geometry_element.attrib["rgba"].split(" ")
                 geometry_element.set("rgba", f"{r} {g} {b} 0")
 
         # -----------------------

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/mujoco/model.py

@@ -73,7 +73,7 @@ class MujocoModelHelper:
             new_hfield = generate_hfield(heightmap, (nrow, ncol))
             model.hfield_data = new_hfield
 
-        return MujocoModelHelper(model=model, data=mj.MjData(model))
+        return MujocoModelHelper(model=model, data=data)
 
     def time(self) -> float:
         """Return the simulation time."""

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/mujoco/visualizer.py

@@ -173,4 +173,4 @@ class MujocoVisualizer:
         try:
             yield handle
         finally:
-            handle.close() if close_on_exit else None
+            _ = handle.close() if close_on_exit else None

{jaxsim-0.2.1.dev128 → jaxsim-0.2.1.dev155}/src/jaxsim/parsers/descriptions/collision.py

@@ -1,11 +1,13 @@
+from __future__ import annotations
+
 import abc
 import dataclasses
-from typing import List
 
 import jax.numpy as jnp
 import numpy as np
 import numpy.typing as npt
 
+import jaxsim.typing as jtp
 from jaxsim import logging
 
 from .link import LinkDescription
@@ -17,9 +19,9 @@ class CollidablePoint:
     Represents a collidable point associated with a parent link.
 
     Attributes:
-        parent_link (LinkDescription): The parent link to which the collidable point is attached.
-        position (npt.NDArray): The position of the collidable point relative to the parent link.
-        enabled (bool): A flag indicating whether the collidable point is enabled for collision detection.
+        parent_link: The parent link to which the collidable point is attached.
+        position: The position of the collidable point relative to the parent link.
+        enabled: A flag indicating whether the collidable point is enabled for collision detection.
 
     """
 
@@ -29,7 +31,7 @@ class CollidablePoint:
 
     def change_link(
         self, new_link: LinkDescription, new_H_old: npt.NDArray
-    ) -> "CollidablePoint":
+    ) -> CollidablePoint:
         """
         Move the collidable point to a new parent link.
 
@@ -39,8 +41,8 @@ class CollidablePoint:
 
         Returns:
             CollidablePoint: A new collidable point associated with the new parent link.
-
         """
+
         msg = f"Moving collidable point: {self.parent_link.name} -> {new_link.name}"
         logging.debug(msg=msg)
 
@@ -50,15 +52,24 @@ class CollidablePoint:
             enabled=self.enabled,
         )
 
-    def __eq__(self, other):
-        retval = (
-            self.parent_link == other.parent_link
-            and (self.position == other.position).all()
-            and self.enabled == other.enabled
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(self.parent_link),
+                hash(tuple(self.position.tolist())),
+                hash(self.enabled),
+            )
         )
-        return retval
 
-    def __str__(self):
+    def __eq__(self, other: CollidablePoint) -> bool:
+
+        if not isinstance(other, CollidablePoint):
+            return False
+
+        return hash(self) == hash(other)
+
+    def __str__(self) -> str:
         return (
             f"{self.__class__.__name__}("
             + f"parent_link={self.parent_link.name}"
@@ -74,11 +85,11 @@ class CollisionShape(abc.ABC):
     Abstract base class for representing collision shapes.
 
     Attributes:
-        collidable_points (List[CollidablePoint]): A list of collidable points associated with the collision shape.
+        collidable_points: A list of collidable points associated with the collision shape.
 
     """
 
-    collidable_points: List[CollidablePoint]
+    collidable_points: tuple[CollidablePoint]
 
     def __str__(self):
         return (
@@ -95,14 +106,26 @@ class BoxCollision(CollisionShape):
     Represents a box-shaped collision shape.
 
     Attributes:
-        center (npt.NDArray): The center of the box in the local frame of the collision shape.
+        center: The center of the box in the local frame of the collision shape.
 
     """
 
-    center: npt.NDArray
+    center: jtp.VectorLike
 
-    def __eq__(self, other):
-        return (self.center == other.center).all() and super().__eq__(other)
+    def __hash__(self) -> int:
+        return hash(
+            (
+                hash(super()),
+                hash(tuple(self.center.tolist())),
+            )
+        )
+
+    def __eq__(self, other: BoxCollision) -> bool:
+
+        if not isinstance(other, BoxCollision):
+            return False
+
+        return hash(self) == hash(other)
 
 
 @dataclasses.dataclass
@@ -111,11 +134,23 @@ class SphereCollision(CollisionShape):
     Represents a spherical collision shape.
 
     Attributes:
-        center (npt.NDArray): The center of the sphere in the local frame of the collision shape.
+        center: The center of the sphere in the local frame of the collision shape.
 
     """
 
-    center: npt.NDArray
+    center: jtp.VectorLike
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                hash(super()),
+                hash(tuple(self.center.tolist())),
+            )
+        )
+
+    def __eq__(self, other: BoxCollision) -> bool:
+
+        if not isinstance(other, BoxCollision):
+            return False
 
-    def __eq__(self, other):
-        return (self.center == other.center).all() and super().__eq__(other)
+        return hash(self) == hash(other)