jaxsim 0.5.1.dev123__py3-none-any.whl → 0.5.1.dev133__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/com.py

@@ -279,7 +279,7 @@ def bias_acceleration(
         C_v̇_WL: jtp.Vector, C_v_WC: jtp.Vector, L_H_C: jtp.Matrix, L_v_LC: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the body-fixed representation of the link bias acceleration
+        Convert the body-fixed representation of the link bias acceleration
         C_v̇_WL expressed in a generic frame C to the body-fixed representation L_v̇_WL.
         """
 

jaxsim/api/common.py

@@ -26,7 +26,7 @@ _R = TypeVar("_R")
 
 
 def named_scope(fn, name: str | None = None) -> Callable[_P, _R]:
-    """Applies a JAX named scope to a function for improved profiling and clarity."""
+    """Apply a JAX named scope to a function for improved profiling and clarity."""
 
     @functools.wraps(fn)
     def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:

jaxsim/api/contact.py

@@ -293,6 +293,9 @@ def in_contact(
 def estimate_good_soft_contacts_parameters(
     *args, **kwargs
 ) -> jaxsim.rbda.contacts.ContactParamsTypes:
+    """
+    Estimate good soft contacts parameters. Deprecated, use `estimate_good_contact_parameters` instead.
+    """
 
     msg = "This method is deprecated, please use `{}`."
     logging.warning(msg.format(estimate_good_contact_parameters.__name__))

jaxsim/api/data.py

@@ -456,7 +456,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     @jax.jit
     def generalized_velocity(self) -> jtp.Vector:
         r"""
-        Get the generalized velocity
+        Get the generalized velocity.
+
         :math:`\boldsymbol{\nu} = (\boldsymbol{v}_{W,B};\, \boldsymbol{\omega}_{W,B};\, \mathbf{s}) \in \mathbb{R}^{6+n}`
 
         Returns:

jaxsim/api/kin_dyn_parameters.py

@@ -16,7 +16,7 @@ from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
 
 
 @jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
-class KynDynParameters(JaxsimDataclass):
+class KinDynParameters(JaxsimDataclass):
     r"""
     Class storing the kinematic and dynamic parameters of a model.
 
@@ -52,14 +52,20 @@ class KynDynParameters(JaxsimDataclass):
 
     @property
     def parent_array(self) -> jtp.Vector:
+        r"""
+        Return the parent array :math:`\lambda(i)` of the model.
+        """
         return self._parent_array.get()
 
     @property
     def support_body_array_bool(self) -> jtp.Matrix:
+        r"""
+        Return the boolean support parent array :math:`\kappa_{b}(i)` of the model.
+        """
         return self._support_body_array_bool.get()
 
     @staticmethod
-    def build(model_description: ModelDescription) -> KynDynParameters:
+    def build(model_description: ModelDescription) -> KinDynParameters:
         """
         Construct the kinematic and dynamic parameters of the model.
 
@@ -210,10 +216,10 @@ class KynDynParameters(JaxsimDataclass):
         )
 
         # =================================
-        # Build and return KynDynParameters
+        # Build and return KinDynParameters
         # =================================
 
-        return KynDynParameters(
+        return KinDynParameters(
             link_names=tuple(l.name for l in ordered_links),
             _parent_array=HashedNumpyArray(array=parent_array),
             _support_body_array_bool=HashedNumpyArray(array=support_body_array_bool),
@@ -224,9 +230,9 @@ class KynDynParameters(JaxsimDataclass):
             frame_parameters=frame_parameters,
         )
 
-    def __eq__(self, other: KynDynParameters) -> bool:
+    def __eq__(self, other: KinDynParameters) -> bool:
 
-        if not isinstance(other, KynDynParameters):
+        if not isinstance(other, KinDynParameters):
             return False
 
         return hash(self) == hash(other)
@@ -450,7 +456,7 @@ class KynDynParameters(JaxsimDataclass):
 
     def set_link_mass(
         self, link_index: jtp.IntLike, mass: jtp.FloatLike
-    ) -> KynDynParameters:
+    ) -> KinDynParameters:
         """
         Set the mass of a link.
 
@@ -470,7 +476,7 @@ class KynDynParameters(JaxsimDataclass):
 
     def set_link_inertia(
         self, link_index: jtp.IntLike, inertia: jtp.MatrixLike
-    ) -> KynDynParameters:
+    ) -> KinDynParameters:
         r"""
         Set the inertia tensor of a link.
 
@@ -648,7 +654,16 @@ class LinkParameters(JaxsimDataclass):
     def build_from_flat_parameters(
         index: jtp.IntLike, parameters: jtp.VectorLike
     ) -> LinkParameters:
+        """
+        Build a LinkParameters object from a flat vector of parameters.
+
+        Args:
+            index: The index of the link.
+            parameters: The flat vector of parameters.
 
+        Returns:
+            The LinkParameters object.
+        """
         index = jnp.array(index).squeeze().astype(int)
 
         m = jnp.array(parameters[0]).squeeze().astype(float)
@@ -772,7 +787,9 @@ class ContactParameters(JaxsimDataclass):
 
     @property
     def indices_of_enabled_collidable_points(self) -> npt.NDArray:
-
+        """
+        Return the indices of the enabled collidable points.
+        """
         return np.where(np.array(self.enabled))[0]
 
     @staticmethod

jaxsim/api/model.py

@@ -45,7 +45,7 @@ class JaxSimModel(JaxsimDataclass):
         default=None, repr=False
     )
 
-    kin_dyn_parameters: js.kin_dyn_parameters.KynDynParameters | None = (
+    kin_dyn_parameters: js.kin_dyn_parameters.KinDynParameters | None = (
         dataclasses.field(default=None, repr=False)
     )
 
@@ -63,6 +63,9 @@ class JaxSimModel(JaxsimDataclass):
 
     @property
     def description(self) -> ModelDescription:
+        """
+        Return the model description.
+        """
         return self._description.get()
 
     def __eq__(self, other: JaxSimModel) -> bool:
@@ -271,7 +274,7 @@ class JaxSimModel(JaxsimDataclass):
         # Build the model.
         model = cls(
             model_name=model_name,
-            kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
+            kin_dyn_parameters=js.kin_dyn_parameters.KinDynParameters.build(
                 model_description=model_description
             ),
             time_step=time_step,
@@ -1015,7 +1018,7 @@ def forward_dynamics_aba(
         W_v̇_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WB: jtp.Vector, W_v_WC: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the inertial-fixed apparent base acceleration W_v̇_WB to
+        Convert the inertial-fixed apparent base acceleration W_v̇_WB to
         another representation C_v̇_WB expressed in a generic frame C.
         """
 
@@ -1376,7 +1379,7 @@ def inverse_dynamics(
 
     def to_inertial(C_v̇_WB, W_H_C, C_v_WB, W_v_WC):
         """
-        Helper to convert the active representation of the base acceleration C_v̇_WB
+        Convert the active representation of the base acceleration C_v̇_WB
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
@@ -1825,7 +1828,7 @@ def link_bias_accelerations(
         C_v̇_WB: jtp.Vector, C_v_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WC: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the active representation of the base acceleration C_v̇_WB
+        Convert the active representation of the base acceleration C_v̇_WB
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
@@ -1961,7 +1964,7 @@ def link_bias_accelerations(
         L_v̇_WL: jtp.Vector, L_v_WL: jtp.Vector, C_H_L: jtp.Matrix, L_v_CL: jtp.Vector
     ) -> jtp.Vector:
         """
-        Helper to convert the body-fixed apparent acceleration L_v̇_WL to
+        Convert the body-fixed apparent acceleration L_v̇_WL to
         another representation C_v̇_WL expressed in a generic frame C.
         """
 

jaxsim/api/ode.py

@@ -15,12 +15,32 @@ from .ode_data import ODEState
 
 
 class SystemDynamicsFromModelAndData(Protocol):
+    """
+    Protocol defining the signature of a function computing the system dynamics
+    given a model and data object.
+    """
+
     def __call__(
         self,
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         **kwargs: dict[str, Any],
-    ) -> tuple[ODEState, dict[str, Any]]: ...
+    ) -> tuple[ODEState, dict[str, Any]]:
+        """
+        Compute the system dynamics given a model and data object.
+
+        Args:
+            model: The model to consider.
+            data: The data of the considered model.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            A tuple with an `ODEState` object storing in each of its attributes the
+            corresponding derivative, and the dictionary of auxiliary data returned
+            by the system dynamics evaluation.
+        """
+
+        pass
 
 
 def wrap_system_dynamics_for_integration(

jaxsim/exceptions.py

@@ -17,6 +17,8 @@ def raise_if(
         msg:
             The message to display when the exception is raised. The message can be a
             format string (fmt), whose fields are filled with the args and kwargs.
+        *args: The arguments to fill the format string.
+        **kwargs: The keyword arguments to fill the format string
     """
 
     # Disable host callback if running on unsupported hardware or if the user
@@ -61,6 +63,9 @@ def raise_if(
 def raise_runtime_error_if(
     condition: bool | jax.Array, msg: str, *args, **kwargs
 ) -> None:
+    """
+    Raise a RuntimeError if a condition is met. Useful in jit-compiled functions.
+    """
 
     return raise_if(condition, RuntimeError, msg, *args, **kwargs)
 
@@ -68,5 +73,8 @@ def raise_runtime_error_if(
 def raise_value_error_if(
     condition: bool | jax.Array, msg: str, *args, **kwargs
 ) -> None:
+    """
+    Raise a ValueError if a condition is met. Useful in jit-compiled functions.
+    """
 
     return raise_if(condition, ValueError, msg, *args, **kwargs)

jaxsim/integrators/common.py

@@ -36,9 +36,25 @@ PyTreeType = TypeVar("PyTreeType", bound=jtp.PyTree)
 
 
 class SystemDynamics(Protocol[State, StateDerivative]):
+    """
+    Protocol defining the system dynamics.
+    """
+
     def __call__(
         self, x: State, t: Time, **kwargs
-    ) -> tuple[StateDerivative, dict[str, Any]]: ...
+    ) -> tuple[StateDerivative, dict[str, Any]]:
+        """
+        Compute the state derivative of the system.
+
+        Args:
+            x: The state of the system.
+            t: The time of the system.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The state derivative of the system and the auxiliary dictionary.
+        """
+        pass
 
 
 # =======================
@@ -48,6 +64,9 @@ class SystemDynamics(Protocol[State, StateDerivative]):
 
 @jax_dataclasses.pytree_dataclass
 class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
+    """
+    Factory class for integrators.
+    """
 
     dynamics: Static[SystemDynamics[State, StateDerivative]] = dataclasses.field(
         repr=False, hash=False, compare=False, kw_only=True
@@ -110,6 +129,9 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
     def __call__(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
     ) -> tuple[NextState, dict[str, Any]]:
+        """
+        Perform a single integration step.
+        """
         pass
 
     def init(
@@ -121,6 +143,9 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
         include_dynamics_aux_dict: bool = False,
         **kwargs,
     ) -> dict[str, Any]:
+        """
+        Initialize the integrator. This method is deprecated.
+        """
 
         logging.warning(
             "The 'init' method has been deprecated. There is no need to call it."
@@ -131,6 +156,18 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
 @jax_dataclasses.pytree_dataclass
 class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]):
+    """
+    Base class for explicit Runge-Kutta integrators.
+
+    Attributes:
+        A: The Runge-Kutta matrix.
+        b: The weights coefficients.
+        c: The nodes coefficients.
+        order_of_bT_rows: The order of the solution.
+        row_index_of_solution: The row of the integration output corresponding to the final solution.
+        fsal_enabled_if_supported: Whether to enable the FSAL property, if supported.
+        index_of_fsal: The index of the intermediate derivative to be used as the first derivative of the next iteration.
+    """
 
     # The Runge-Kutta matrix.
     A: ClassVar[jtp.Matrix]
@@ -156,10 +193,16 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
     @property
     def has_fsal(self) -> bool:
+        """
+        Check if the integrator supports the FSAL property.
+        """
         return self.fsal_enabled_if_supported and self.index_of_fsal is not None
 
     @property
     def order(self) -> int:
+        """
+        Return the order of the integrator.
+        """
         return self.order_of_bT_rows[self.row_index_of_solution]
 
     @override
@@ -221,6 +264,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
     def __call__(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
     ) -> tuple[NextState, dict[str, Any]]:
+        """
+        Perform a single integration step.
+        """
 
         # Here z is a batched state with as many batch elements as b.T rows.
         # Note that z has multiple batches only if b.T has more than one row,
@@ -331,7 +377,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         def scan_body(
             carry: jax.Array, i: int | jax.Array
         ) -> tuple[jax.Array, dict[str, Any]]:
-            """"""
+            """
+            Compute the kᵢ derivative of the Runge-Kutta stage.
+            """
 
             # Unpack the carry, i.e. the stacked kᵢ vectors.
             K = carry
@@ -498,6 +546,16 @@ class ExplicitRungeKuttaSO3Mixin:
     def post_process_state(
         cls, x0: js.ode_data.ODEState, t0: Time, xf: js.ode_data.ODEState, dt: TimeStep
     ) -> js.ode_data.ODEState:
+        r"""
+        Post-process the integrated state at :math:`t_f = t_0 + \Delta t` so that the
+        quaternion is normalized.
+
+        Args:
+            x0: The initial state of the system.
+            t0: The initial time of the system.
+            xf: The final state of the system obtain through the integration.
+            dt: The time step used for the integration.
+        """
 
         # Extract the initial base quaternion.
         W_Q_B_t0 = x0.physics_model.base_quaternion

jaxsim/integrators/fixed_step.py

@@ -17,6 +17,9 @@ ODEStateDerivative = js.ode_data.ODEState
 
 @jax_dataclasses.pytree_dataclass
 class ForwardEuler(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    Forward Euler integrator.
+    """
 
     A: ClassVar[jtp.Matrix] = jnp.atleast_2d(0).astype(float)
 
@@ -30,6 +33,9 @@ class ForwardEuler(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
 @jax_dataclasses.pytree_dataclass
 class Heun2(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    Heun's second-order integrator.
+    """
 
     A: ClassVar[jtp.Matrix] = jnp.array(
         [
@@ -56,6 +62,9 @@ class Heun2(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
 @jax_dataclasses.pytree_dataclass
 class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    Fourth-order Runge-Kutta integrator.
+    """
 
     A: ClassVar[jtp.Matrix] = jnp.array(
         [
@@ -89,14 +98,26 @@ class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
 @jax_dataclasses.pytree_dataclass
 class ForwardEulerSO3(ExplicitRungeKuttaSO3Mixin, ForwardEuler[js.ode_data.ODEState]):
+    """
+    Forward Euler integrator for SO(3) states.
+    """
+
     pass
 
 
 @jax_dataclasses.pytree_dataclass
 class Heun2SO3(ExplicitRungeKuttaSO3Mixin, Heun2[js.ode_data.ODEState]):
+    """
+    Heun's second-order integrator for SO(3) states.
+    """
+
     pass
 
 
 @jax_dataclasses.pytree_dataclass
 class RungeKutta4SO3(ExplicitRungeKuttaSO3Mixin, RungeKutta4[js.ode_data.ODEState]):
+    """
+    Fourth-order Runge-Kutta integrator for SO(3) states.
+    """
+
     pass

jaxsim/integrators/variable_step.py

@@ -216,6 +216,17 @@ def local_error_estimation(
 
 @jax_dataclasses.pytree_dataclass
 class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
+    """
+    An Embedded Runge-Kutta integrator.
+
+    This class implements a general-purpose Embedded Runge-Kutta integrator
+    that can be used to solve ordinary differential equations with adaptive
+    step sizes.
+
+    The integrator is based on an Explicit Runge-Kutta method, and it uses
+    two different solutions to estimate the local integration error. The
+    error is then used to adapt the step size to reach a desired accuracy.
+    """
 
     AfterInitKey: ClassVar[str] = "after_init"
     InitializingKey: ClassVar[str] = "initializing"
@@ -257,6 +268,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
             x0: The initial state of the system.
             t0: The initial time of the system.
             dt: The time step of the integration.
+            **kwargs: Additional parameters.
 
         Returns:
             The metadata of the integrator to be passed to the first step.
@@ -296,6 +308,9 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
     def __call__(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
     ) -> tuple[NextState, dict[str, Any]]:
+        """
+        Integrate the system for a single step.
+        """
 
         # This method is called differently in three stages:
         #
@@ -512,10 +527,16 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
     @property
     def order_of_solution(self) -> int:
+        """
+        The order of the solution.
+        """
         return self.order_of_bT_rows[self.row_index_of_solution]
 
     @property
     def order_of_solution_estimate(self) -> int:
+        """
+        The order of the solution estimate.
+        """
         return self.order_of_bT_rows[self.row_index_of_solution_estimate]
 
     @classmethod
@@ -534,6 +555,23 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         max_step_rejections: jtp.IntLike = MAX_STEP_REJECTIONS_DEFAULT,
         **kwargs,
     ) -> Self:
+        """
+        Build an Embedded Runge-Kutta integrator.
+
+        Args:
+            dynamics: The system dynamics function.
+            fsal_enabled_if_supported:
+                Whether to enable the FSAL property if supported by the integrator.
+            dt_max: The maximum step size.
+            dt_min: The minimum step size.
+            rtol: The relative tolerance.
+            atol: The absolute tolerance.
+            safety: The safety factor to shrink the step size.
+            beta_max: The maximum factor to increase the step size.
+            beta_min: The minimum factor to increase the step size.
+            max_step_rejections: The maximum number of step rejections.
+            **kwargs: Additional parameters.
+        """
 
         # Check that b.T has enough rows based on the configured index of the
         # solution estimate. This is necessary for embedded methods.
@@ -569,6 +607,9 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
 @jax_dataclasses.pytree_dataclass
 class HeunEulerSO3(EmbeddedRungeKutta[PyTreeType], ExplicitRungeKuttaSO3Mixin):
+    """
+    The Heun-Euler integrator for SO(3) dynamics.
+    """
 
     A: ClassVar[jtp.Matrix] = jnp.array(
         [
@@ -602,6 +643,9 @@ class HeunEulerSO3(EmbeddedRungeKutta[PyTreeType], ExplicitRungeKuttaSO3Mixin):
 
 @jax_dataclasses.pytree_dataclass
 class BogackiShampineSO3(EmbeddedRungeKutta[PyTreeType], ExplicitRungeKuttaSO3Mixin):
+    """
+    The Bogacki-Shampine integrator for SO(3) dynamics.
+    """
 
     A: ClassVar[jtp.Matrix] = jnp.array(
         [

jaxsim/math/adjoint.py

@@ -7,6 +7,10 @@ from .skew import Skew
 
 
 class Adjoint:
+    """
+    A utility class for adjoint matrix operations.
+    """
+
     @staticmethod
     def from_quaternion_and_translation(
         quaternion: jtp.Vector = jnp.array([1.0, 0, 0, 0]),
@@ -18,11 +22,10 @@ class Adjoint:
         Create an adjoint matrix from a quaternion and a translation.
 
         Args:
-            quaternion (jtp.Vector): A quaternion vector (4D) representing orientation.
-            translation (jtp.Vector): A translation vector (3D).
-            inverse (bool): Whether to compute the inverse adjoint. Default is False.
-            normalize_quaternion (bool): Whether to normalize the quaternion before creating the adjoint.
-                                         Default is False.
+            quaternion: A quaternion vector (4D) representing orientation.
+            translation: A translation vector (3D).
+            inverse: Whether to compute the inverse adjoint.
+            normalize_quaternion: Whether to normalize the quaternion before creating the adjoint.
 
         Returns:
             jtp.Matrix: The adjoint matrix.
@@ -69,9 +72,9 @@ class Adjoint:
         Create an adjoint matrix from a rotation matrix and a translation vector.
 
         Args:
-            rotation (jtp.Matrix): A 3x3 rotation matrix.
-            translation (jtp.Vector): A translation vector (3D).
-            inverse (bool): Whether to compute the inverse adjoint. Default is False.
+            rotation: A 3x3 rotation matrix.
+            translation: A translation vector (3D).
+            inverse: Whether to compute the inverse adjoint. Default is False.
 
         Returns:
             jtp.Matrix: The adjoint matrix.
@@ -105,7 +108,7 @@ class Adjoint:
         Convert an adjoint matrix to a transformation matrix.
 
         Args:
-            adjoint (jtp.Matrix): The adjoint matrix (6x6).
+            adjoint: The adjoint matrix (6x6).
 
         Returns:
             jtp.Matrix: The transformation matrix (4x4).
@@ -131,7 +134,7 @@ class Adjoint:
         Compute the inverse of an adjoint matrix.
 
         Args:
-            adjoint (jtp.Matrix): The adjoint matrix.
+            adjoint: The adjoint matrix.
 
         Returns:
             jtp.Matrix: The inverse adjoint matrix.

jaxsim/math/cross.py

@@ -6,13 +6,17 @@ from .skew import Skew
 
 
 class Cross:
+    """
+    A utility class for cross product matrix operations.
+    """
+
     @staticmethod
     def vx(velocity_sixd: jtp.Vector) -> jtp.Matrix:
         """
         Compute the cross product matrix for 6D velocities.
 
         Args:
-            velocity_sixd (jtp.Vector): A 6D velocity vector [v, ω].
+            velocity_sixd: A 6D velocity vector [v, ω].
 
         Returns:
             jtp.Matrix: The cross product matrix (6x6).
@@ -37,7 +41,7 @@ class Cross:
         Compute the negative transpose of the cross product matrix for 6D velocities.
 
         Args:
-            velocity_sixd (jtp.Vector): A 6D velocity vector [v, ω].
+            velocity_sixd: A 6D velocity vector [v, ω].
 
         Returns:
             jtp.Matrix: The negative transpose of the cross product matrix (6x6).

jaxsim/math/inertia.py

@@ -6,15 +6,19 @@ from .skew import Skew
 
 
 class Inertia:
+    """
+    A utility class for inertia matrix operations.
+    """
+
     @staticmethod
     def to_sixd(mass: jtp.Float, com: jtp.Vector, I: jtp.Matrix) -> jtp.Matrix:
         """
         Convert mass, center of mass, and inertia matrix to a 6x6 inertia matrix.
 
         Args:
-            mass (jtp.Float): The mass of the body.
-            com (jtp.Vector): The center of mass position (3D).
-            I (jtp.Matrix): The 3x3 inertia matrix.
+            mass: The mass of the body.
+            com: The center of mass position (3D).
+            I: The 3x3 inertia matrix.
 
         Returns:
             jtp.Matrix: The 6x6 inertia matrix.
@@ -42,7 +46,7 @@ class Inertia:
         Convert a 6x6 inertia matrix to mass, center of mass, and inertia matrix.
 
         Args:
-            M (jtp.Matrix): The 6x6 inertia matrix.
+            M: The 6x6 inertia matrix.
 
         Returns:
             tuple[jtp.Float, jtp.Vector, jtp.Matrix]: A tuple containing mass, center of mass (3D), and inertia matrix (3x3).