jaxsim 0.2.1.dev123__py3-none-any.whl → 0.2.1.dev155__py3-none-any.whl

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/_version.py

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

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/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/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/api/model.py

@@ -16,7 +16,8 @@ from jax_dataclasses import Static
 import jaxsim.api as js
 import jaxsim.parsers.descriptions
 import jaxsim.typing as jtp
-from jaxsim.utils import HashlessObject, JaxsimDataclass, Mutability
+from jaxsim.math import Cross
+from jaxsim.utils import JaxsimDataclass, Mutability
 
 from .common import VelRepr
 
@@ -32,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)
     )
@@ -40,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:
 
@@ -60,6 +58,7 @@ class JaxSimModel(JaxsimDataclass):
         return hash(
             (
                 hash(self.model_name),
+                hash(self.description),
                 hash(self.kin_dyn_parameters),
             )
         )
@@ -156,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
             ),
@@ -301,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.
     """
 
@@ -314,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
@@ -685,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()
@@ -871,6 +867,126 @@ def free_floating_mass_matrix(
             raise ValueError(data.velocity_representation)
 
 
+@jax.jit
+def free_floating_coriolis_matrix(
+    model: JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    """
+    Compute the free-floating Coriolis matrix of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The free-floating Coriolis matrix of the model.
+
+    Note:
+        This function, contrarily to other quantities of the equations of motion,
+        does not exploit any iterative algorithm. Therefore, the computation of
+        the Coriolis matrix may be much slower than other quantities.
+    """
+
+    # We perform all the calculation in body-fixed.
+    # The Coriolis matrix computed in this representation is converted later
+    # to the active representation stored in data.
+    with data.switch_velocity_representation(VelRepr.Body):
+
+        B_ν = data.generalized_velocity()
+
+        # Doubly-left free-floating Jacobian.
+        L_J_WL_B = generalized_free_floating_jacobian(model=model, data=data)
+
+        # Doubly-left free-floating Jacobian derivative.
+        L_J̇_WL_B = jax.vmap(
+            lambda link_index: js.link.jacobian_derivative(
+                model=model, data=data, link_index=link_index
+            )
+        )(js.link.names_to_idxs(model=model, link_names=model.link_names()))
+
+    L_M_L = link_spatial_inertia_matrices(model=model)
+
+    # Body-fixed link velocities.
+    # Note: we could have called link.velocity() instead of computing it ourselves,
+    # but since we need the link Jacobians later, we can save a double calculation.
+    L_v_WL = jax.vmap(lambda J: J @ B_ν)(L_J_WL_B)
+
+    # Compute the contribution of each link to the Coriolis matrix.
+    def compute_link_contribution(M, v, J, J̇) -> jtp.Array:
+
+        return J.T @ ((Cross.vx_star(v) @ M + M @ Cross.vx(v)) @ J + M @ J̇)
+
+    C_B_links = jax.vmap(compute_link_contribution)(
+        L_M_L,
+        L_v_WL,
+        L_J_WL_B,
+        L_J̇_WL_B,
+    )
+
+    # We need to adjust the Coriolis matrix for fixed-base models.
+    # In this case, the base link does not contribute to the matrix, and we need to zero
+    # the off-diagonal terms mapping joint quantities onto the base configuration.
+    if model.floating_base():
+        C_B = C_B_links.sum(axis=0)
+    else:
+        C_B = C_B_links[1:].sum(axis=0)
+        C_B = C_B.at[0:6, 6:].set(0.0)
+        C_B = C_B.at[6:, 0:6].set(0.0)
+
+    # Adjust the representation of the Coriolis matrix.
+    # Refer to https://github.com/traversaro/traversaro-phd-thesis, Section 3.6.
+    match data.velocity_representation:
+
+        case VelRepr.Body:
+            return C_B
+
+        case VelRepr.Inertial:
+
+            n = model.dofs()
+            W_H_B = data.base_transform()
+            B_X_W = jaxsim.math.Adjoint.from_transform(W_H_B, inverse=True)
+            B_T_W = jax.scipy.linalg.block_diag(B_X_W, jnp.eye(n))
+
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                W_v_WB = data.base_velocity()
+                B_Ẋ_W = -B_X_W @ jaxsim.math.Cross.vx(W_v_WB)
+
+            B_Ṫ_W = jax.scipy.linalg.block_diag(B_Ẋ_W, jnp.zeros(shape=(n, n)))
+
+            with data.switch_velocity_representation(VelRepr.Body):
+                M = free_floating_mass_matrix(model=model, data=data)
+
+            C = B_T_W.T @ (M @ B_Ṫ_W + C_B @ B_T_W)
+
+            return C
+
+        case VelRepr.Mixed:
+
+            n = model.dofs()
+            BW_H_B = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
+            B_X_BW = jaxsim.math.Adjoint.from_transform(transform=BW_H_B, inverse=True)
+            B_T_BW = jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(n))
+
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                BW_v_WB = data.base_velocity()
+                BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
+
+            BW_v_BW_B = BW_v_WB - BW_v_W_BW
+            B_Ẋ_BW = -B_X_BW @ jaxsim.math.Cross.vx(BW_v_BW_B)
+
+            B_Ṫ_BW = jax.scipy.linalg.block_diag(B_Ẋ_BW, jnp.zeros(shape=(n, n)))
+
+            with data.switch_velocity_representation(VelRepr.Body):
+                M = free_floating_mass_matrix(model=model, data=data)
+
+            C = B_T_BW.T @ (M @ B_Ṫ_BW + C_B @ B_T_BW)
+
+            return C
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+
 @jax.jit
 def inverse_dynamics(
     model: JaxSimModel,
@@ -931,8 +1047,6 @@ def inverse_dynamics(
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
-        from jaxsim.math import Cross
-
         W_X_C = jaxlie.SE3.from_matrix(W_H_C).adjoint()
         C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
         C_v_WC = C_X_W @ W_v_WC
@@ -1364,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
@@ -1410,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/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/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/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/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/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/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)

jaxsim/parsers/descriptions/joint.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import ClassVar, Tuple, Union
+from typing import ClassVar
 
 import jax_dataclasses
 import numpy as np
-import numpy.typing as npt
 
 import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass, Mutability
@@ -15,6 +14,7 @@ from .link import LinkDescription
 
 @dataclasses.dataclass(frozen=True)
 class JointType:
+
     Fixed: ClassVar[int] = 0
     Revolute: ClassVar[int] = 1
     Prismatic: ClassVar[int] = 2
@@ -64,29 +64,31 @@ class JointDescription(JaxsimDataclass):
     """
 
     name: jax_dataclasses.Static[str]
-    axis: npt.NDArray
-    pose: npt.NDArray
-    jtype: jax_dataclasses.Static[JointType]
+    axis: jtp.Vector
+    pose: jtp.Matrix
+    jtype: jax_dataclasses.Static[jtp.IntLike]
     child: LinkDescription = dataclasses.dataclass(repr=False)
     parent: LinkDescription = dataclasses.dataclass(repr=False)
 
-    index: int | None = None
+    index: jtp.IntLike | None = None
+
+    friction_static: jtp.FloatLike = 0.0
+    friction_viscous: jtp.FloatLike = 0.0
 
-    friction_static: float = 0.0
-    friction_viscous: float = 0.0
+    position_limit_damper: jtp.FloatLike = 0.0
+    position_limit_spring: jtp.FloatLike = 0.0
 
-    position_limit_damper: float = 0.0
-    position_limit_spring: float = 0.0
+    position_limit: tuple[jtp.FloatLike, jtp.FloatLike] = (0.0, 0.0)
+    initial_position: jtp.FloatLike | jtp.VectorLike = 0.0
 
-    position_limit: Tuple[float, float] = (0.0, 0.0)
-    initial_position: Union[float, npt.NDArray] = 0.0
+    motor_inertia: jtp.FloatLike = 0.0
+    motor_viscous_friction: jtp.FloatLike = 0.0
+    motor_gear_ratio: jtp.FloatLike = 1.0
 
-    motor_inertia: float = 0.0
-    motor_viscous_friction: float = 0.0
-    motor_gear_ratio: float = 1.0
+    def __post_init__(self) -> None:
 
-    def __post_init__(self):
         if self.axis is not None:
+
             with self.mutable_context(
                 mutability=Mutability.MUTABLE, restore_after_exception=False
             ):
@@ -94,4 +96,24 @@ class JointDescription(JaxsimDataclass):
                 self.axis = self.axis / norm_of_axis
 
     def __hash__(self) -> int:
-        return hash(self.__repr__())
+
+        return hash(
+            (
+                hash(self.name),
+                hash(tuple(self.axis.tolist())),
+                hash(tuple(self.pose.flatten().tolist())),
+                hash(int(self.jtype)),
+                hash(self.child),
+                hash(self.parent),
+                hash(int(self.index)) if self.index is not None else 0,
+                hash(float(self.friction_static)),
+                hash(float(self.friction_viscous)),
+                hash(float(self.position_limit_damper)),
+                hash(float(self.position_limit_spring)),
+                hash((float(el) for el in self.position_limit)),
+                hash(tuple(np.atleast_1d(self.initial_position).tolist())),
+                hash(float(self.motor_inertia)),
+                hash(float(self.motor_viscous_friction)),
+                hash(float(self.motor_gear_ratio)),
+            ),
+        )

jaxsim/parsers/descriptions/link.py

@@ -5,6 +5,7 @@ import dataclasses
 import jax.numpy as jnp
 import jax_dataclasses
 import jaxlie
+import numpy as np
 from jax_dataclasses import Static
 
 import jaxsim.typing as jtp
@@ -23,7 +24,7 @@ class LinkDescription(JaxsimDataclass):
         index: An optional index for the link (it gets automatically assigned).
         parent: The parent link of this link.
         pose: The pose transformation matrix of the link.
-        children: List of child links.
+        children: The children links.
     """
 
     name: Static[str]
@@ -33,7 +34,7 @@ class LinkDescription(JaxsimDataclass):
     parent: LinkDescription = dataclasses.field(default=None, repr=False)
     pose: jtp.Matrix = dataclasses.field(default_factory=lambda: jnp.eye(4), repr=False)
 
-    children: Static[list[LinkDescription]] = dataclasses.field(
+    children: Static[tuple[LinkDescription]] = dataclasses.field(
         default_factory=list, repr=False
     )
 
@@ -43,10 +44,12 @@ class LinkDescription(JaxsimDataclass):
             (
                 hash(self.name),
                 hash(float(self.mass)),
-                hash(tuple(self.inertia.flatten().tolist())),
-                hash(int(self.index)),
-                hash(self.parent),
-                hash(tuple(hash(c) for c in self.children)),
+                hash(tuple(np.atleast_1d(self.inertia).flatten().tolist())),
+                hash(int(self.index)) if self.index is not None else 0,
+                hash(tuple(np.atleast_1d(self.pose).flatten().tolist())),
+                hash(tuple(self.children)),
+                # Here only using the name to prevent circular recursion:
+                hash(self.parent.name) if self.parent is not None else 0,
             )
         )
 

jaxsim/parsers/descriptions/model.py

@@ -27,7 +27,7 @@ class ModelDescription(KinematicGraph):
 
     fixed_base: bool = True
 
-    collision_shapes: list[CollisionShape] = dataclasses.field(
+    collision_shapes: tuple[CollisionShape, ...] = dataclasses.field(
         default_factory=list, repr=False, hash=False
     )
 
@@ -37,7 +37,7 @@ class ModelDescription(KinematicGraph):
         links: list[LinkDescription],
         joints: list[JointDescription],
         frames: list[LinkDescription] | None = None,
-        collisions: list[CollisionShape] = (),
+        collisions: tuple[CollisionShape, ...] = (),
         fixed_base: bool = False,
         base_link_name: str | None = None,
         considered_joints: Sequence[str] | None = None,
@@ -87,7 +87,7 @@ class ModelDescription(KinematicGraph):
         for collision_shape in collisions:
 
             # Get all the collidable points of the shape
-            coll_points = list(collision_shape.collidable_points)
+            coll_points = tuple(collision_shape.collidable_points)
 
             # Assume they have an unique parent link
             if not len(set({cp.parent_link.name for cp in coll_points})) == 1:
@@ -111,7 +111,7 @@ class ModelDescription(KinematicGraph):
                 continue
 
             # Create a new collision shape
-            new_collision_shape = CollisionShape(collidable_points=[])
+            new_collision_shape = CollisionShape(collidable_points=())
             final_collisions.append(new_collision_shape)
 
             # If the frame was found, update the collidable points' pose and add them
@@ -133,19 +133,19 @@ class ModelDescription(KinematicGraph):
                     ),
                 )
 
-                # Store the updated collision
-                new_collision_shape.collidable_points.append(moved_cp)
+                # Store the updated collision.
+                new_collision_shape.collidable_points += (moved_cp,)
 
         # Build the model
         model = ModelDescription(
             name=name,
             root_pose=kinematic_graph.root_pose,
             fixed_base=fixed_base,
-            collision_shapes=final_collisions,
+            collision_shapes=tuple(final_collisions),
             root=kinematic_graph.root,
             joints=kinematic_graph.joints,
             frames=kinematic_graph.frames,
-            _joints_removed=kinematic_graph._joints_removed,
+            _joints_removed=kinematic_graph.joints_removed,
         )
 
         # Check that the root link of kinematic graph is the desired base link.
@@ -174,7 +174,7 @@ class ModelDescription(KinematicGraph):
             links=list(self.links_dict.values()),
             joints=self.joints,
             frames=self.frames,
-            collisions=self.collision_shapes,
+            collisions=tuple(self.collision_shapes),
             fixed_base=self.fixed_base,
             base_link_name=list(iter(self))[0].name,
             model_pose=self.root_pose,
@@ -182,8 +182,8 @@ class ModelDescription(KinematicGraph):
         )
 
         # Include the unconnected/removed joints from the original model.
-        for joint in self._joints_removed:
-            reduced_model_description._joints_removed.append(joint)
+        for joint in self.joints_removed:
+            reduced_model_description.joints_removed.append(joint)
 
         return reduced_model_description
 
@@ -243,3 +243,23 @@ class ModelDescription(KinematicGraph):
 
         # Return enabled collidable points
         return [cp for cp in all_collidable_points if cp.enabled]
+
+    def __eq__(self, other: ModelDescription) -> bool:
+
+        if not isinstance(other, ModelDescription):
+            return False
+
+        return hash(self) == hash(other)
+
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(self.name),
+                hash(self.fixed_base),
+                hash(self.root),
+                hash(tuple(self.joints)),
+                hash(tuple(self.frames)),
+                hash(self.root_pose),
+            )
+        )

jaxsim/parsers/kinematic_graph.py

@@ -31,14 +31,21 @@ class RootPose(NamedTuple):
     root_position: npt.NDArray = np.zeros(3)
     root_quaternion: npt.NDArray = np.array([1.0, 0, 0, 0])
 
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(tuple(self.root_position.tolist())),
+                hash(tuple(self.root_quaternion.tolist())),
+            )
+        )
+
     def __eq__(self, other: RootPose) -> bool:
 
         if not isinstance(other, RootPose):
             return False
 
-        return np.allclose(self.root_position, other.root_position) and np.allclose(
-            self.root_quaternion, other.root_quaternion
-        )
+        return hash(self) == hash(other)
 
 
 @dataclasses.dataclass(frozen=True)
@@ -54,22 +61,24 @@ class KinematicGraph(Sequence[descriptions.LinkDescription]):
     """
 
     root: descriptions.LinkDescription
-    frames: list[descriptions.LinkDescription] = dataclasses.field(default_factory=list)
+    frames: list[descriptions.LinkDescription] = dataclasses.field(
+        default_factory=list, hash=False, compare=False
+    )
     joints: list[descriptions.JointDescription] = dataclasses.field(
-        default_factory=list
+        default_factory=list, hash=False, compare=False
     )
 
     root_pose: RootPose = dataclasses.field(default_factory=lambda: RootPose())
 
     # Private attribute storing optional additional info.
     _extra_info: dict[str, Any] = dataclasses.field(
-        repr=False, compare=False, default_factory=dict
+        default_factory=dict, repr=False, hash=False, compare=False
     )
 
     # Private attribute storing the unconnected joints from the parsed model and
     # the joints removed after model reduction.
     _joints_removed: list[descriptions.JointDescription] = dataclasses.field(
-        default_factory=list, repr=False, compare=False
+        default_factory=list, repr=False, hash=False, compare=False
     )
 
     @functools.cached_property
@@ -98,14 +107,17 @@ class KinematicGraph(Sequence[descriptions.LinkDescription]):
         for index, link in enumerate(self):
             link.mutable(validate=False).index = index
 
-        # Get the names of the links and frames.
+        # Get the names of the links, frames, and joints.
         link_names = [l.name for l in self]
         frame_names = [f.name for f in self.frames]
+        joint_names = [j.name for j in self.joints]
 
         # Make sure that they are unique.
         assert len(link_names) == len(set(link_names))
         assert len(frame_names) == len(set(frame_names))
+        assert len(joint_names) == len(set(joint_names))
         assert set(link_names).isdisjoint(set(frame_names))
+        assert set(link_names).isdisjoint(set(joint_names))
 
         # Order frames with their name.
         super().__setattr__("frames", sorted(self.frames, key=lambda f: f.name))
@@ -251,7 +263,7 @@ class KinematicGraph(Sequence[descriptions.LinkDescription]):
 
         # Reset the connections of the root link.
         for link in links_dict.values():
-            link.children = []
+            link.children = tuple()
 
         # Couple links and joints creating the kinematic graph.
         for joint in joints:
@@ -268,7 +280,8 @@ class KinematicGraph(Sequence[descriptions.LinkDescription]):
 
             # Assign link's children and make sure they are unique.
             if child_link.name not in {l.name for l in parent_link.children}:
-                parent_link.children.append(child_link)
+                with parent_link.mutable_context(Mutability.MUTABLE_NO_VALIDATION):
+                    parent_link.children = parent_link.children + (child_link,)
 
         # Collect all the links of the kinematic graph.
         all_links_in_graph = list(
@@ -315,7 +328,7 @@ class KinematicGraph(Sequence[descriptions.LinkDescription]):
         # Update the unconnected links by removing their children. The other properties
         # are left untouched, it's caller responsibility to post-process them if needed.
         for link in unconnected_links:
-            link.children = []
+            link.children = tuple()
             msg = "Link '{}' won't be part of the kinematic graph because unconnected"
             logging.debug(msg=msg.format(link.name))
 
@@ -615,6 +628,17 @@ class KinematicGraph(Sequence[descriptions.LinkDescription]):
             horizontal=True,
         )
 
+    @property
+    def joints_removed(self) -> list[descriptions.JointDescription]:
+        """
+        Get the list of joints removed during the graph reduction.
+
+        Returns:
+            The list of removed joints.
+        """
+
+        return self._joints_removed
+
     @staticmethod
     def breadth_first_search(
         root: descriptions.LinkDescription,
@@ -785,6 +809,7 @@ class KinematicGraphTransforms:
 
             # Get the joint.
             joint = self.graph.joints_dict[name]
+            assert joint.name == name
 
             # Get the transform of the parent link.
             M_H_L = self.transform(name=joint.parent.name)

jaxsim/parsers/rod/parser.py

@@ -7,7 +7,7 @@ import numpy as np
 import rod
 
 from jaxsim import logging
-from jaxsim.math.quaternion import Quaternion
+from jaxsim.math import Quaternion
 from jaxsim.parsers import descriptions, kinematic_graph
 
 from . import utils

jaxsim/parsers/rod/utils.py

@@ -6,7 +6,7 @@ import numpy.typing as npt
 import rod
 
 import jaxsim.typing as jtp
-from jaxsim.math.inertia import Inertia
+from jaxsim.math import Inertia
 from jaxsim.parsers import descriptions
 
 
@@ -59,9 +59,7 @@ def from_sdf_inertial(inertial: rod.Inertial) -> jtp.Matrix:
     return M_L.astype(dtype=float)
 
 
-def joint_to_joint_type(
-    joint: rod.Joint,
-) -> descriptions.JointType:
+def joint_to_joint_type(joint: rod.Joint) -> int:
     """
     Extract the joint type from an SDF joint.
 
@@ -69,7 +67,7 @@ def joint_to_joint_type(
         joint: The parsed SDF joint.
 
     Returns:
-        The corresponding joint type description.
+        The integer corresponding to the joint type.
     """
 
     axis = joint.axis
@@ -138,7 +136,7 @@ def create_box_collision(
     collidable_points = [
         descriptions.CollidablePoint(
             parent_link=link_description,
-            position=corner,
+            position=np.array(corner),
             enabled=True,
         )
         for corner in box_corners_wrt_link.T
@@ -197,7 +195,7 @@ def create_sphere_collision(
     collidable_points = [
         descriptions.CollidablePoint(
             parent_link=link_description,
-            position=point,
+            position=np.array(point),
             enabled=True,
         )
         for point in sphere_points_wrt_link.T

jaxsim/rbda/crba.py

@@ -111,7 +111,7 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
         # a while loop using a for loop with fixed number of iterations.
         def inner_fn(carry: CarryInnerFn, k: jtp.Int) -> tuple[CarryInnerFn, None]:
             def compute_inner(carry: CarryInnerFn) -> tuple[CarryInnerFn, None]:
-                j, Fi, M = carry
+                j, _, _ = carry
                 out = jax.lax.cond(
                     pred=(λ[j] > 0),
                     true_fun=while_loop_body,
@@ -120,7 +120,7 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
                 )
                 return out, None
 
-            j, Fi, M = carry
+            j, _, _ = carry
             return jax.lax.cond(
                 pred=(k == j),
                 true_fun=compute_inner,

jaxsim/rbda/forward_kinematics.py

@@ -49,7 +49,7 @@ def forward_kinematics_model(
     # Compute the parent-to-child adjoints and the motion subspaces of the joints.
     # These transforms define the relative kinematics of the entire model, including
     # the base transform for both floating-base and fixed-base models.
-    i_X_λi, S = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
+    i_X_λi, _ = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
         joint_positions=s, base_transform=W_H_B.as_matrix()
     )
 

jaxsim/typing.py

@@ -1,4 +1,5 @@
-from typing import Any, Hashable
+from collections.abc import Hashable
+from typing import Any
 
 import jax
 
@@ -24,6 +25,7 @@ PyTree = (
 # =======================
 
 Array = jax.typing.ArrayLike
+Scalar = Array
 Vector = Array
 Matrix = Array
 
@@ -31,6 +33,7 @@ Int = int | IntJax
 Bool = bool | ArrayJax
 Float = float | FloatJax
 
+ScalarLike = Scalar | int | float
 ArrayLike = Array
 VectorLike = Vector
 MatrixLike = Matrix

{jaxsim-0.2.1.dev123.dist-info → jaxsim-0.2.1.dev155.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev123
+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 == 'viz'
 
 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.dev123.dist-info → jaxsim-0.2.1.dev155.dist-info}/RECORD

@@ -1,18 +1,18 @@
-jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=jAmAP2ZYBVoFw5W8EyAYOEpIKED5Jy8XK259JnSs0Fk,428
+jaxsim/__init__.py,sha256=xzuTuZrgKdWLqqDzbvqzm2cJrEtAbepOeUqDu7ByVek,2621
+jaxsim/_version.py,sha256=VTwovSkbhBoA3d5Ku6UfV56QTou0p3bVvIXRd63p0yQ,428
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
-jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
+jaxsim/typing.py,sha256=cl7HHQCeP3mHmtF6EuQZcCjGvDmc_AryMWntP_lRBGg,722
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=Yof6otFi-mLWAs1rqjmeNJTOWIH9gn7BdU5EIjiL6Ts,13481
 jaxsim/api/common.py,sha256=bqQ__pIQZbh-j8rkoHUkYHAgGiJnDzjHG-q4Ny0OOYQ,6646
-jaxsim/api/contact.py,sha256=Cvr-EfQtHP3nymtWdo-9WWU24Bkta-2Pp3nKsdjo6uc,12778
+jaxsim/api/contact.py,sha256=79kcdq7C1_kWgxd1QWBabBhIPkwWEVLk-Fiz9kh-4so,12800
 jaxsim/api/data.py,sha256=xfKJz6Rw0YTk-EHCGiT8BFQrs_ggOz01lRi1Qh1mb28,27256
 jaxsim/api/frame.py,sha256=0YXOrGmx3cSQqa4_Ky-n6zyup3I3xvXNEgub-Bc5xUw,6222
 jaxsim/api/joint.py,sha256=-5DogPg4g4mmLckyVIVNjwv-Rxz0IWS7_md9nDlhPWA,4581
-jaxsim/api/kin_dyn_parameters.py,sha256=zMca7OmCsCWK_cavLTSZSeYh9Qu1-409cdsyWvWPAUQ,26090
-jaxsim/api/link.py,sha256=oW5-DShmmeCRk3JOJIwzo3HbWuNGmpm_wBJ4fkmrROM,16645
-jaxsim/api/model.py,sha256=1HlQ5FMzeJAk-cE1pmELgVjzMYUX9-iipw3N4WssAL4,55435
-jaxsim/api/ode.py,sha256=BfvV_14uu0szWecoDiV8rTu-dvSFLK7eyrO38ZqHB_w,10157
+jaxsim/api/kin_dyn_parameters.py,sha256=b1e96I8hKU5fh4StLObdVcDpr_6ZglrgD3SRyrqTu18,26203
+jaxsim/api/link.py,sha256=MdMWaMpM5Dj5JHK8uwHZ4zR4Fjq3R4asi2sGTxk1OAs,16647
+jaxsim/api/model.py,sha256=sCx9CcP23A1I_ae4UqTq4Fpq5u0aDki72CqgnR1H50w,59465
+jaxsim/api/ode.py,sha256=luTQJsIXUtCp_81dR42X7WrMvwrXtYbyJiqss29v7zA,10786
 jaxsim/api/ode_data.py,sha256=D6FzMkvY_qNuoFEImyp7sxAk-0pJOd3oZeSr9bBTcLk,23089
 jaxsim/api/references.py,sha256=UA6kSQVBoq-bXSo99EOELf-_MD5MTy2zS0GtG3wQ410,16618
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
@@ -23,31 +23,31 @@ jaxsim/math/__init__.py,sha256=inJ9nRFkqstuGa8OyFkfWVudo5U9Ug4WgDBuKva8AIA,337
 jaxsim/math/adjoint.py,sha256=DT21izjVW497GRrgNfx8tv0ZeWW5QncWMGMhI0acUNw,4425
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
 jaxsim/math/inertia.py,sha256=UAB7ym4gXFanejcs_ovZMpteHCc6poWYmt-mLmd5hhk,1640
-jaxsim/math/joint_model.py,sha256=hfkEQOglk265P9W0-nvXxI95k21x2deYEv0mud6XUKc,9989
+jaxsim/math/joint_model.py,sha256=xJSocGOyLzLJIQo4j5rBfMCPD4ltUQ2jCZfN747i2Ck,9989
 jaxsim/math/quaternion.py,sha256=X9b8jHf0QemKUjIZSnXRJc3DdMr42CBhBy_mi9_X_AM,5068
 jaxsim/math/rotation.py,sha256=Z90daUjGpuNEVLfWB3SVtM9EtwAIaneVj9A9UpWXqhA,2182
 jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=ZKzoXwKRowE17p3P0EgKrq8mocdMTGUKNyaKt6-bAX0,2941
 jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=Tq3b2tq_BulGm9GCM00NQMzcD7B0yKvoyiP94aFbjx4,21023
-jaxsim/mujoco/model.py,sha256=5-4KTbEbU19zjrSuvUVdLo3noWxTvlCNsFIs3rQTNDY,13506
-jaxsim/mujoco/visualizer.py,sha256=YlteqcCbeB1B6saAHKBz1IJad3N5Rp163reZrzKLAys,5065
+jaxsim/mujoco/loaders.py,sha256=7rjpeJ6_GuitlCty-ZkLhTILQ0GmsFzDMgve-7Gkkh4,20984
+jaxsim/mujoco/model.py,sha256=1KVRjSLOTCuHt53apBPQTnFYJRknlVoKLQaxWsNK8qc,13494
+jaxsim/mujoco/visualizer.py,sha256=PXgQzwetS9mRJYHBknDMLsQ9152FdrSvZuT9xE_dfIQ,5069
 jaxsim/parsers/__init__.py,sha256=sonYi-bBWAoB04kp1mxT4uIORxjb7SdZ0ukGPmVx98Y,44
-jaxsim/parsers/kinematic_graph.py,sha256=zFt7x7pPGJar36Azukdi1eI_sa1kMWD3B8kZqcHx6iw,33934
+jaxsim/parsers/kinematic_graph.py,sha256=WdIxntWfxXf67x90oM5KHHXFrSITMwVahqWgcOjYFzc,34730
 jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
-jaxsim/parsers/descriptions/collision.py,sha256=HUWwuRgI9KznY29FFw1_zU3bGigDEezrcPOJSxSJGNU,3382
-jaxsim/parsers/descriptions/joint.py,sha256=lRnYMmjpASpz0Ueuqzwnj5Ze4yLRgPTx66H0_kbQnNI,3042
-jaxsim/parsers/descriptions/link.py,sha256=GC-6ZgRZuRVpcRo1sY6YaR8lkCHkR4DvHNs2Ydw_tn4,2887
-jaxsim/parsers/descriptions/model.py,sha256=uO5xOJtViihVPnSSsmfQJvCh45ANyi9KYAzLOhH0R8g,8993
+jaxsim/parsers/descriptions/collision.py,sha256=BQeIG-TKi4SVny23w6riDrQ5itC6VRwEMBX6HgAXHxA,3973
+jaxsim/parsers/descriptions/joint.py,sha256=z_nYSS0fdkcaerjUlPX0U1Vn1ArBT0u_XdKjqxG3HcY,3959
+jaxsim/parsers/descriptions/link.py,sha256=QvEE7J6iMQLibpLqlcBV428UA7NMpFFXJwe35GYnjAY,3124
+jaxsim/parsers/descriptions/model.py,sha256=V9nSyCK3mo7680WYMDEx1MTfdDTJzbCGPqAp3qA2XRE,9511
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=Q13TOkmpU0SHpgSV8WRYWb290aPNNLsaz4eMlD4Mq5w,13525
-jaxsim/parsers/rod/utils.py,sha256=9oO4YsQRaR2v700IkNOXRnPpn5i4N8HFfgjPkMLK2mc,5732
+jaxsim/parsers/rod/parser.py,sha256=4COuhkAYv4-GIpCqvkXEJWpDEQczEkBM3KwpqX48Rek,13514
+jaxsim/parsers/rod/utils.py,sha256=KSjgy6WsmTrD5HZEA2x8hOBSRU4bUGOOHzxKkeFO5r8,5721
 jaxsim/rbda/__init__.py,sha256=MqEZwzu8SHPAlIFHmSXmCjehuOJGRX58OrBVAbBVMwg,374
 jaxsim/rbda/aba.py,sha256=0OoCzHhf1v-qqr1y5PIrD7_mPwAlid0fjXxUrIa5E_s,9118
 jaxsim/rbda/collidable_points.py,sha256=4ZNJbEj2nEi15jBLR-GNbdaqKgkN58FBgqd_TXupEgg,4948
-jaxsim/rbda/crba.py,sha256=GodskOZjtrSlbQAqxRv1un_706O7BaJK-U2qa18vJk8,4741
-jaxsim/rbda/forward_kinematics.py,sha256=OHugNU7C0UxYAW0o1rqH1ZgniSwurz6L1T1MJxfxq08,3418
+jaxsim/rbda/crba.py,sha256=awsWEQXLE0UPEXIcZCVsAqBEPjyahMNzY9ux6nE1l-s,4739
+jaxsim/rbda/forward_kinematics.py,sha256=94W7TUXvZjMb-99CyYR8pObuxIYYX9B_dtRZqsNcThs,3418
 jaxsim/rbda/jacobian.py,sha256=M79bGir-2w_iJ2GurYhOGgMfJnp7ZMOCW6AeeWKK8iM,10745
 jaxsim/rbda/rnea.py,sha256=DjwkvXQVUSUclM3Uy3UPZ2tao91R5dGd4o7TsS2qObI,7650
 jaxsim/rbda/soft_contacts.py,sha256=52zJOF31hFpqoaOednTvi8j_UxhRcdGNjzOPb2v2MPc,11257
@@ -58,8 +58,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=EJMcblYKUjxw9HJShVf81Ig3pHUJno6Dx6h-RnY--wM,2040
-jaxsim-0.2.1.dev123.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.2.1.dev123.dist-info/METADATA,sha256=hOUnUjhHwGvdZVMR_-z5pKFsd681swunNdYkqgC338I,9745
-jaxsim-0.2.1.dev123.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.1.dev123.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.1.dev123.dist-info/RECORD,,
+jaxsim-0.2.1.dev155.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.2.1.dev155.dist-info/METADATA,sha256=Jhyuk3qGnK7WzHYbuA4b5mCXofQw7WZAj6zqh01jKkU,9740
+jaxsim-0.2.1.dev155.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.1.dev155.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.1.dev155.dist-info/RECORD,,