jaxsim 0.3.1.dev62__py3-none-any.whl → 0.3.1.dev94__py3-none-any.whl

jaxsim/api/model.py

@@ -9,14 +9,14 @@ from typing import Any, Sequence
 import jax
 import jax.numpy as jnp
 import jax_dataclasses
-import jaxlie
 import rod
 from jax_dataclasses import Static
 
 import jaxsim.api as js
-import jaxsim.parsers.descriptions
+import jaxsim.terrain
 import jaxsim.typing as jtp
-from jaxsim.math import Cross
+from jaxsim.math import Adjoint, Cross
+from jaxsim.parsers.descriptions import ModelDescription
 from jaxsim.utils import JaxsimDataclass, Mutability, wrappers
 
 from .common import VelRepr
@@ -46,12 +46,12 @@ class JaxSimModel(JaxsimDataclass):
         default=None, repr=False
     )
 
-    _description: Static[
-        wrappers.HashlessObject[jaxsim.parsers.descriptions.ModelDescription | None]
-    ] = dataclasses.field(default=None, repr=False)
+    _description: Static[wrappers.HashlessObject[ModelDescription | None]] = (
+        dataclasses.field(default=None, repr=False)
+    )
 
     @property
-    def description(self) -> jaxsim.parsers.descriptions.ModelDescription:
+    def description(self) -> ModelDescription:
         return self._description.get()
 
     def __eq__(self, other: JaxSimModel) -> bool:
@@ -116,7 +116,7 @@ class JaxSimModel(JaxsimDataclass):
         import jaxsim.parsers.rod
 
         # Parse the input resource (either a path to file or a string with the URDF/SDF)
-        # and build the -intermediate- model description
+        # and build the -intermediate- model description.
         intermediate_description = jaxsim.parsers.rod.build_model_description(
             model_description=model_description, is_urdf=is_urdf
         )
@@ -128,7 +128,7 @@ class JaxSimModel(JaxsimDataclass):
                 considered_joints=considered_joints
             )
 
-        # Build the model
+        # Build the model.
         model = JaxSimModel.build(
             model_description=intermediate_description,
             model_name=model_name,
@@ -136,7 +136,7 @@ class JaxSimModel(JaxsimDataclass):
             contact_model=contact_model,
         )
 
-        # Store the origin of the model, in case downstream logic needs it
+        # Store the origin of the model, in case downstream logic needs it.
         with model.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
             model.built_from = model_description
 
@@ -144,7 +144,7 @@ class JaxSimModel(JaxsimDataclass):
 
     @staticmethod
     def build(
-        model_description: jaxsim.parsers.descriptions.ModelDescription,
+        model_description: ModelDescription,
         model_name: str | None = None,
         *,
         terrain: jaxsim.terrain.Terrain | None = None,
@@ -169,14 +169,14 @@ class JaxSimModel(JaxsimDataclass):
         """
         from jaxsim.rbda.contacts.soft import SoftContacts
 
-        # Set the model name (if not provided, use the one from the model description)
+        # Set the model name (if not provided, use the one from the model description).
         model_name = model_name if model_name is not None else model_description.name
 
-        # Set the terrain (if not provided, use the default flat terrain)
+        # Set the terrain (if not provided, use the default flat terrain).
         terrain = terrain or JaxSimModel.__dataclass_fields__["terrain"].default
         contact_model = contact_model or SoftContacts(terrain=terrain)
 
-        # Build the model
+        # Build the model.
         model = JaxSimModel(
             model_name=model_name,
             _description=wrappers.HashlessObject(obj=model_description),
@@ -361,7 +361,7 @@ def reduce(
         considered_joints=list(considered_joints)
     )
 
-    # Build the reduced model
+    # Build the reduced model.
     reduced_model = JaxSimModel.build(
         model_description=reduced_intermediate_description,
         model_name=model.name(),
@@ -369,7 +369,7 @@ def reduce(
         contact_model=model.contact_model,
     )
 
-    # Store the origin of the model, in case downstream logic needs it
+    # Store the origin of the model, in case downstream logic needs it.
     with reduced_model.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
         reduced_model.built_from = model.built_from
 
@@ -493,7 +493,7 @@ def generalized_free_floating_jacobian(
         case VelRepr.Inertial:
 
             W_H_B = data.base_transform()
-            B_X_W = jaxlie.SE3.from_matrix(W_H_B).inverse().adjoint()
+            B_X_W = Adjoint.from_transform(transform=W_H_B, inverse=True)
 
             B_J_full_WX_I = B_J_full_WX_W = B_J_full_WX_B @ jax.scipy.linalg.block_diag(
                 B_X_W, jnp.eye(model.dofs())
@@ -507,7 +507,7 @@ def generalized_free_floating_jacobian(
 
             W_R_B = data.base_orientation(dcm=True)
             BW_H_B = jnp.eye(4).at[0:3, 0:3].set(W_R_B)
-            B_X_BW = jaxlie.SE3.from_matrix(BW_H_B).inverse().adjoint()
+            B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
 
             B_J_full_WX_I = B_J_full_WX_BW = (
                 B_J_full_WX_B
@@ -715,7 +715,7 @@ def forward_dynamics_aba(
 
         # 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()
+        C_X_W = Adjoint.from_transform(transform=W_H_C, inverse=True)
         return C_X_W @ (W_v̇_WB - Cross.vx(W_v_WC) @ W_v_WB)
 
     match data.velocity_representation:
@@ -797,21 +797,21 @@ def forward_dynamics_crb(
     # Prepare data
     # ============
 
-    # Build joint torques if not provided
+    # Build joint torques if not provided.
     τ = (
         jnp.atleast_1d(joint_forces)
         if joint_forces is not None
         else jnp.zeros_like(data.joint_positions())
     )
 
-    # Build external forces if not provided
+    # Build external forces if not provided.
     f = (
         jnp.atleast_2d(link_forces)
         if link_forces is not None
         else jnp.zeros(shape=(model.number_of_links(), 6))
     )
 
-    # Compute terms of the floating-base EoM
+    # Compute terms of the floating-base EoM.
     M = free_floating_mass_matrix(model=model, data=data)
     h = free_floating_bias_forces(model=model, data=data)
     S = jnp.block([jnp.zeros(shape=(model.dofs(), 6)), jnp.eye(model.dofs())]).T
@@ -848,7 +848,7 @@ def forward_dynamics_crb(
     # 6D transformation X.
     v̇_WB = ν̇[0:6].squeeze().astype(float)
 
-    # Extract the joint accelerations
+    # Extract the joint accelerations.
     s̈ = jnp.atleast_1d(ν̇[6:].squeeze()).astype(float)
 
     return v̇_WB, s̈
@@ -880,7 +880,9 @@ def free_floating_mass_matrix(
 
         case VelRepr.Inertial:
 
-            B_X_W = jaxlie.SE3.from_matrix(data.base_transform()).inverse().adjoint()
+            B_X_W = Adjoint.from_transform(
+                transform=data.base_transform(), inverse=True
+            )
             invT = jax.scipy.linalg.block_diag(B_X_W, jnp.eye(model.dofs()))
 
             return invT.T @ M_body @ invT
@@ -888,7 +890,7 @@ def free_floating_mass_matrix(
         case VelRepr.Mixed:
 
             BW_H_B = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
-            B_X_BW = jaxlie.SE3.from_matrix(BW_H_B).inverse().adjoint()
+            B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
             invT = jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
 
             return invT.T @ M_body @ invT
@@ -1077,8 +1079,8 @@ def inverse_dynamics(
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
-        W_X_C = jaxlie.SE3.from_matrix(W_H_C).adjoint()
-        C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
+        W_X_C = Adjoint.from_transform(transform=W_H_C)
+        C_X_W = Adjoint.from_transform(transform=W_H_C, inverse=True)
         C_v_WC = C_X_W @ W_v_WC
 
         # In Mixed representation, we need to include a cross product in ℝ⁶.
@@ -1187,12 +1189,12 @@ def free_floating_gravity_forces(
         The free-floating gravity forces of the model.
     """
 
-    # Build a zeroed state
+    # Build a zeroed state.
     data_rnea = js.data.JaxSimModelData.zero(
         model=model, velocity_representation=data.velocity_representation
     )
 
-    # Set just the generalized position
+    # Set just the generalized position.
     with data_rnea.mutable_context(
         mutability=Mutability.MUTABLE, restore_after_exception=False
     ):
@@ -1237,12 +1239,12 @@ def free_floating_bias_forces(
         The free-floating bias forces of the model.
     """
 
-    # Build a zeroed state
+    # Build a zeroed state.
     data_rnea = js.data.JaxSimModelData.zero(
         model=model, velocity_representation=data.velocity_representation
     )
 
-    # Set the generalized position and generalized velocity
+    # Set the generalized position and generalized velocity.
     with data_rnea.mutable_context(
         mutability=Mutability.MUTABLE, restore_after_exception=False
     ):
@@ -1361,12 +1363,14 @@ def total_momentum_jacobian(
             B_Jh = B_Jh_B
 
         case VelRepr.Inertial:
-            B_X_W = jaxlie.SE3.from_matrix(data.base_transform()).inverse().adjoint()
+            B_X_W = Adjoint.from_transform(
+                transform=data.base_transform(), inverse=True
+            )
             B_Jh = B_Jh_B @ jax.scipy.linalg.block_diag(B_X_W, jnp.eye(model.dofs()))
 
         case VelRepr.Mixed:
             BW_H_B = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
-            B_X_BW = jaxlie.SE3.from_matrix(BW_H_B).inverse().adjoint()
+            B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
             B_Jh = B_Jh_B @ jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
 
         case _:
@@ -1378,14 +1382,14 @@ def total_momentum_jacobian(
 
         case VelRepr.Inertial:
             W_H_B = data.base_transform()
-            B_Xv_W = jaxlie.SE3.from_matrix(W_H_B).inverse().adjoint()
+            B_Xv_W = Adjoint.from_transform(transform=W_H_B, inverse=True)
             W_Xf_B = B_Xv_W.T
             W_Jh = W_Xf_B @ B_Jh
             return W_Jh
 
         case VelRepr.Mixed:
             BW_H_B = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
-            B_Xv_BW = jaxlie.SE3.from_matrix(BW_H_B).inverse().adjoint()
+            B_Xv_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
             BW_Xf_B = B_Xv_BW.T
             BW_Jh = BW_Xf_B @ B_Jh
             return BW_Jh
@@ -1449,7 +1453,7 @@ def average_velocity_jacobian(
             W_p_CoM = js.com.com_position(model=model, data=data)
 
             W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
-            W_X_GW = jaxlie.SE3.from_matrix(W_H_GW).adjoint()
+            W_X_GW = Adjoint.from_transform(transform=W_H_GW)
 
             return W_X_GW @ GW_J
 
@@ -1461,7 +1465,7 @@ def average_velocity_jacobian(
             B_R_W = data.base_orientation(dcm=True).transpose()
 
             B_H_GB = jnp.eye(4).at[0:3, 3].set(B_R_W @ (W_p_CoM - W_p_B))
-            B_X_GB = jaxlie.SE3.from_matrix(B_H_GB).adjoint()
+            B_X_GB = Adjoint.from_transform(transform=B_H_GB)
 
             return B_X_GB @ GB_J
 
@@ -1472,7 +1476,7 @@ def average_velocity_jacobian(
             W_p_CoM = js.com.com_position(model=model, data=data)
 
             BW_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM - W_p_B)
-            BW_X_GW = jaxlie.SE3.from_matrix(BW_H_GW).adjoint()
+            BW_X_GW = Adjoint.from_transform(transform=BW_H_GW)
 
             return BW_X_GW @ GW_J
 
@@ -1518,8 +1522,8 @@ def link_bias_accelerations(
         expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
         """
 
-        W_X_C = jaxlie.SE3.from_matrix(W_H_C).adjoint()
-        C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
+        W_X_C = Adjoint.from_transform(transform=W_H_C)
+        C_X_W = Adjoint.from_transform(transform=W_H_C, inverse=True)
 
         # In Mixed representation, we need to include a cross product in ℝ⁶.
         # In Inertial and Body representations, the cross product is always zero.
@@ -1606,7 +1610,7 @@ def link_bias_accelerations(
     # not remove gravity during the propagation.
 
     # Initialize the loop.
-    Carry = tuple[jtp.MatrixJax, jtp.MatrixJax]
+    Carry = tuple[jtp.Matrix, jtp.Matrix]
     carry0: Carry = (L_v_WL, L_v̇_WL)
 
     def propagate_accelerations(carry: Carry, i: jtp.Int) -> tuple[Carry, None]:
@@ -1832,8 +1836,8 @@ def step(
         integrator_state: The state of the integrator.
         joint_forces: The joint forces to consider.
         link_forces:
-            The link 6D forces to consider.
-            The frame in which they are expressed must be `data.velocity_representation`.
+            The 6D forces to apply to the links expressed in the frame corresponding to
+            the velocity representation of `data`.
         kwargs: Additional kwargs to pass to the integrator.
 
     Returns:

jaxsim/api/ode.py

@@ -96,7 +96,9 @@ def system_velocity_dynamics(
         model: The model to consider.
         data: The data of the considered model.
         joint_forces: The joint forces to apply.
-        link_forces: The 6D forces to apply to the links.
+        link_forces:
+            The 6D forces to apply to the links expressed in the frame corresponding to
+            the velocity representation of `data`.
 
     Returns:
         A tuple containing the derivative of the base 6D velocity in inertial-fixed
@@ -105,14 +107,16 @@ def system_velocity_dynamics(
         the system dynamics evaluation.
     """
 
-    # Build joint torques if not provided
+    # Build joint torques if not provided.
     τ = (
         jnp.atleast_1d(joint_forces.squeeze())
         if joint_forces is not None
         else jnp.zeros_like(data.joint_positions())
     ).astype(float)
 
-    # Build link forces if not provided
+    # Build link forces if not provided.
+    # These forces are expressed in the frame corresponding to the velocity
+    # representation of data.
     O_f_L = (
         jnp.atleast_2d(link_forces.squeeze())
         if link_forces is not None
@@ -127,16 +131,17 @@ def system_velocity_dynamics(
     # with the terrain.
     W_f_Li_terrain = jnp.zeros_like(O_f_L).astype(float)
 
-    # Initialize the 6D contact forces W_f ∈ ℝ^{n_c × 6} applied to collidable points,
-    # expressed in the world frame.
-    W_f_Ci = None
+    # Import privately the soft contacts classes.
+    from jaxsim.rbda.contacts.soft import SoftContactsState
 
     # Initialize the derivative of the tangential deformation ṁ ∈ ℝ^{n_c × 3}.
+    assert isinstance(data.state.contact, SoftContactsState)
     ṁ = jnp.zeros_like(data.state.contact.tangential_deformation).astype(float)
 
     if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
-        # Compute the 6D forces applied to each collidable point and the
-        # corresponding material deformation rates.
+
+        # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
+        #  and the corresponding material deformation rates.
         with data.switch_velocity_representation(VelRepr.Inertial):
             W_f_Ci, ṁ = js.contact.collidable_point_dynamics(model=model, data=data)
 
@@ -178,7 +183,7 @@ def system_velocity_dynamics(
             model.kin_dyn_parameters.joint_parameters.friction_viscous
         ).astype(float)
 
-        # Compute the joint friction torque
+        # Compute the joint friction torque.
         τ_friction = -(
             jnp.diag(kc) @ jnp.sign(data.state.physics_model.joint_velocities)
             + jnp.diag(kv) @ data.state.physics_model.joint_velocities
@@ -188,7 +193,7 @@ def system_velocity_dynamics(
     # Compute forward dynamics
     # ========================
 
-    # Compute the total joint forces
+    # Compute the total joint forces.
     τ_total = τ + τ_friction + τ_position_limit
 
     references = js.references.JaxSimModelReferences.build(
@@ -202,7 +207,7 @@ def system_velocity_dynamics(
     with references.switch_velocity_representation(VelRepr.Inertial):
         W_f_L = references.link_forces(model=model, data=data)
 
-    # Compute the total external 6D forces applied to the links
+    # Compute the total external 6D forces applied to the links.
     W_f_L_total = W_f_L + W_f_Li_terrain
 
     # - Joint accelerations: s̈ ∈ ℝⁿ
@@ -273,7 +278,9 @@ def system_dynamics(
         model: The model to consider.
         data: The data of the considered model.
         joint_forces: The joint forces to apply.
-        link_forces: The 6D forces to apply to the links.
+        link_forces:
+            The 6D forces to apply to the links expressed in the frame corresponding to
+            the velocity representation of `data`.
         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).

jaxsim/api/ode_data.py

@@ -31,8 +31,8 @@ class ODEInput(JaxsimDataclass):
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,
-        joint_forces: jtp.VectorJax | None = None,
-        link_forces: jtp.MatrixJax | None = None,
+        joint_forces: jtp.VectorLike | None = None,
+        link_forces: jtp.MatrixLike | None = None,
     ) -> ODEInput:
         """
         Build an `ODEInput` from a `JaxSimModel`.
@@ -160,7 +160,7 @@ class ODEState(JaxsimDataclass):
             `JaxSimModel` and initialized to zero.
         """
 
-        # Get the contact model from the `JaxSimModel`
+        # Get the contact model from the `JaxSimModel`.
         match model.contact_model:
             case SoftContacts():
                 contact = SoftContactsState.build_from_jaxsim_model(
@@ -212,7 +212,7 @@ class ODEState(JaxsimDataclass):
             else PhysicsModelState.zero(model=model)
         )
 
-        # Get the contact model from the `JaxSimModel`
+        # Get the contact model from the `JaxSimModel`.
         match contact:
             case SoftContactsState():
                 pass
@@ -423,7 +423,7 @@ class PhysicsModelState(JaxsimDataclass):
             base_angular_velocity=jnp.array(base_angular_velocity, dtype=float),
         )
 
-        # assert state.valid(physics_model)
+        # TODO (diegoferigo): assert state.valid(physics_model)
         return physics_model_state
 
     @staticmethod
@@ -501,14 +501,14 @@ class PhysicsModelInput(JaxsimDataclass):
         f_ext: The matrix of external forces applied to the links.
     """
 
-    tau: jtp.VectorJax
-    f_ext: jtp.MatrixJax
+    tau: jtp.Vector
+    f_ext: jtp.Matrix
 
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,
-        joint_forces: jtp.VectorJax | None = None,
-        link_forces: jtp.MatrixJax | None = None,
+        joint_forces: jtp.VectorLike | None = None,
+        link_forces: jtp.MatrixLike | None = None,
     ) -> PhysicsModelInput:
         """
         Build a `PhysicsModelInput` from a `JaxSimModel`.
@@ -535,8 +535,8 @@ class PhysicsModelInput(JaxsimDataclass):
 
     @staticmethod
     def build(
-        joint_forces: jtp.VectorJax | None = None,
-        link_forces: jtp.MatrixJax | None = None,
+        joint_forces: jtp.VectorLike | None = None,
+        link_forces: jtp.MatrixLike | None = None,
         number_of_dofs: jtp.Int | None = None,
         number_of_links: jtp.Int | None = None,
     ) -> PhysicsModelInput:

jaxsim/integrators/common.py

@@ -10,7 +10,6 @@ from jax_dataclasses import Static
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.math import Quaternion
 from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
 
 try:
@@ -28,8 +27,8 @@ except ImportError:
 # Generic types
 # =============
 
-Time = jax.typing.ArrayLike
-TimeStep = jax.typing.ArrayLike
+Time = jtp.FloatLike
+TimeStep = jtp.FloatLike
 State = NextState = TypeVar("State")
 StateDerivative = TypeVar("StateDerivative")
 PyTreeType = TypeVar("PyTreeType", bound=jtp.PyTree)
@@ -80,7 +79,7 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
             The integrator object.
         """
 
-        return cls(dynamics=dynamics, **kwargs)  # noqa
+        return cls(dynamics=dynamics, **kwargs)
 
     def step(
         self,
@@ -192,14 +191,14 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]):
 
     # The Runge-Kutta matrix.
-    A: ClassVar[jax.typing.ArrayLike]
+    A: ClassVar[jtp.Matrix]
 
     # The weights coefficients.
     # Note that in practice we typically use its transpose `b.transpose()`.
-    b: ClassVar[jax.typing.ArrayLike]
+    b: ClassVar[jtp.Matrix]
 
     # The nodes coefficients.
-    c: ClassVar[jax.typing.ArrayLike]
+    c: ClassVar[jtp.Vector]
 
     # Define the order of the solution.
     # It should have as many elements as the number of rows of `b.transpose()`.
@@ -385,7 +384,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             # Define the computation of the Runge-Kutta stage.
             def compute_ki() -> jax.Array:
 
-                # Compute ∑ⱼ aᵢⱼ kⱼ
+                # Compute ∑ⱼ aᵢⱼ kⱼ.
                 op_sum_ak = lambda k: jnp.einsum("s,s...->...", A[i], k)
                 sum_ak = jax.tree_util.tree_map(op_sum_ak, K)
 
@@ -441,7 +440,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
     @staticmethod
     def butcher_tableau_is_valid(
-        A: jax.typing.ArrayLike, b: jax.typing.ArrayLike, c: jax.typing.ArrayLike
+        A: jtp.Matrix, b: jtp.Matrix, c: jtp.Vector
     ) -> jtp.Bool:
         """
         Check if the Butcher tableau is valid.
@@ -467,7 +466,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         return valid
 
     @staticmethod
-    def butcher_tableau_is_explicit(A: jax.typing.ArrayLike) -> jtp.Bool:
+    def butcher_tableau_is_explicit(A: jtp.Matrix) -> jtp.Bool:
         """
         Check if the Butcher tableau corresponds to an explicit integration scheme.
 
@@ -482,9 +481,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
     @staticmethod
     def butcher_tableau_supports_fsal(
-        A: jax.typing.ArrayLike,
-        b: jax.typing.ArrayLike,
-        c: jax.typing.ArrayLike,
+        A: jtp.Matrix,
+        b: jtp.Matrix,
+        c: jtp.Vector,
         index_of_solution: jtp.IntLike = 0,
     ) -> [bool, int | None]:
         """
@@ -548,17 +547,11 @@ class ExplicitRungeKuttaSO3Mixin:
         op = lambda x0_leaf, k_leaf: x0_leaf + dt * k_leaf
         xf: js.ode_data.ODEState = jax.tree_util.tree_map(op, x0, k)
 
-        W_Q_B_t0 = x0.physics_model.base_quaternion
-        W_ω_WB_t0 = x0.physics_model.base_angular_velocity
+        W_Q_B_tf = xf.physics_model.base_quaternion
 
         return xf.replace(
             physics_model=xf.physics_model.replace(
-                base_quaternion=Quaternion.integration(
-                    quaternion=W_Q_B_t0,
-                    dt=dt,
-                    omega=W_ω_WB_t0,
-                    omega_in_body_fixed=False,
-                ),
+                base_quaternion=W_Q_B_tf / jnp.linalg.norm(W_Q_B_tf)
             )
         )
 
@@ -569,10 +562,9 @@ class ExplicitRungeKuttaSO3Mixin:
 
         # Indices to convert quaternions between serializations.
         to_xyzw = jnp.array([1, 2, 3, 0])
-        to_wxyz = jnp.array([3, 0, 1, 2])
 
-        # Get the initial quaternion.
-        W_Q_B_t0 = jaxlie.SO3.from_quaternion_xyzw(
+        # Get the initial rotation.
+        W_R_B_t0 = jaxlie.SO3.from_quaternion_xyzw(
             xyzw=x0.physics_model.base_quaternion[to_xyzw]
         )
 
@@ -582,15 +574,13 @@ class ExplicitRungeKuttaSO3Mixin:
         # on the SO(3) manifold.
         W_ω_WB_tf = xf.physics_model.base_angular_velocity
 
-        # Integrate the quaternion on SO(3).
+        # Integrate the orientation on SO(3).
         # Note that we left-multiply with the exponential map since the angular
         # velocity is expressed in the inertial frame.
-        W_Q_B_tf = jaxlie.SO3.exp(tangent=dt * W_ω_WB_tf) @ W_Q_B_t0
+        W_R_B_tf = jaxlie.SO3.exp(tangent=dt * W_ω_WB_tf) @ W_R_B_t0
 
         # Replace the quaternion in the final state.
         return xf.replace(
-            physics_model=xf.physics_model.replace(
-                base_quaternion=W_Q_B_tf.as_quaternion_xyzw()[to_wxyz]
-            ),
+            physics_model=xf.physics_model.replace(base_quaternion=W_R_B_tf.wxyz),
             validate=True,
         )

jaxsim/integrators/fixed_step.py

@@ -1,10 +1,10 @@
 from typing import ClassVar, Generic
 
-import jax
 import jax.numpy as jnp
 import jax_dataclasses
 
 import jaxsim.api as js
+import jaxsim.typing as jtp
 
 from .common import ExplicitRungeKutta, ExplicitRungeKuttaSO3Mixin, PyTreeType
 
@@ -18,11 +18,11 @@ ODEStateDerivative = js.ode_data.ODEState
 @jax_dataclasses.pytree_dataclass
 class ForwardEuler(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
-    A: ClassVar[jax.typing.ArrayLike] = jnp.atleast_2d(0).astype(float)
+    A: ClassVar[jtp.Matrix] = jnp.atleast_2d(0).astype(float)
 
-    b: ClassVar[jax.typing.ArrayLike] = jnp.atleast_2d(1).astype(float).transpose()
+    b: ClassVar[jtp.Matrix] = jnp.atleast_2d(1).astype(float).transpose()
 
-    c: ClassVar[jax.typing.ArrayLike] = jnp.atleast_1d(0).astype(float)
+    c: ClassVar[jtp.Vector] = jnp.atleast_1d(0).astype(float)
 
     row_index_of_solution: ClassVar[int] = 0
     order_of_bT_rows: ClassVar[tuple[int, ...]] = (1,)
@@ -31,14 +31,14 @@ class ForwardEuler(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 @jax_dataclasses.pytree_dataclass
 class Heun2(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
-    A: ClassVar[jax.typing.ArrayLike] = jnp.array(
+    A: ClassVar[jtp.Matrix] = jnp.array(
         [
             [0, 0],
             [1, 0],
         ]
     ).astype(float)
 
-    b: ClassVar[jax.typing.ArrayLike] = (
+    b: ClassVar[jtp.Matrix] = (
         jnp.atleast_2d(
             jnp.array([1 / 2, 1 / 2]),
         )
@@ -46,7 +46,7 @@ class Heun2(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         .transpose()
     )
 
-    c: ClassVar[jax.typing.ArrayLike] = jnp.array(
+    c: ClassVar[jtp.Vector] = jnp.array(
         [0, 1],
     ).astype(float)
 
@@ -57,7 +57,7 @@ class Heun2(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 @jax_dataclasses.pytree_dataclass
 class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
-    A: ClassVar[jax.typing.ArrayLike] = jnp.array(
+    A: ClassVar[jtp.Matrix] = jnp.array(
         [
             [0, 0, 0, 0],
             [1 / 2, 0, 0, 0],
@@ -66,7 +66,7 @@ class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         ]
     ).astype(float)
 
-    b: ClassVar[jax.typing.ArrayLike] = (
+    b: ClassVar[jtp.Matrix] = (
         jnp.atleast_2d(
             jnp.array([1 / 6, 1 / 3, 1 / 3, 1 / 6]),
         )
@@ -74,7 +74,7 @@ class RungeKutta4(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
         .transpose()
     )
 
-    c: ClassVar[jax.typing.ArrayLike] = jnp.array(
+    c: ClassVar[jtp.Vector] = jnp.array(
         [0, 1 / 2, 1 / 2, 1],
     ).astype(float)