jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev5__py3-none-any.whl

jaxsim/api/model.py

@@ -1,30 +1,30 @@
 from __future__ import annotations
 
+import copy
 import dataclasses
 import functools
 import pathlib
+from collections.abc import Sequence
 from typing import Any
 
 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.physics.algos.aba
-import jaxsim.physics.algos.crba
-import jaxsim.physics.algos.forward_kinematics
-import jaxsim.physics.algos.rnea
-import jaxsim.physics.model.physics_model
+import jaxsim.exceptions
+import jaxsim.terrain
 import jaxsim.typing as jtp
-from jaxsim.high_level.common import VelRepr
-from jaxsim.physics.algos.terrain import FlatTerrain, Terrain
-from jaxsim.utils import JaxsimDataclass, Mutability
+from jaxsim.math import Adjoint, Cross
+from jaxsim.parsers.descriptions import ModelDescription
+from jaxsim.utils import JaxsimDataclass, Mutability, wrappers
 
+from .common import VelRepr
 
-@jax_dataclasses.pytree_dataclass
+
+@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
 class JaxSimModel(JaxsimDataclass):
     """
     The JaxSim model defining the kinematics and dynamics of a robot.
@@ -32,46 +32,87 @@ class JaxSimModel(JaxsimDataclass):
 
     model_name: Static[str]
 
-    physics_model: jaxsim.physics.model.physics_model.PhysicsModel = dataclasses.field(
-        repr=False
+    time_step: jaxsim.integrators.TimeStep = dataclasses.field(
+        default_factory=lambda: jnp.array(0.001, dtype=float),
+    )
+
+    terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
+        default_factory=jaxsim.terrain.FlatTerrain.build, repr=False
+    )
+
+    # Note that this is the default contact model.
+    contact_model: Static[jaxsim.rbda.contacts.ContactModel | None] = dataclasses.field(
+        default=None, repr=False
     )
 
-    terrain: Static[Terrain] = dataclasses.field(default=FlatTerrain(), repr=False)
+    kin_dyn_parameters: js.kin_dyn_parameters.KinDynParameters | None = (
+        dataclasses.field(default=None, repr=False)
+    )
 
     built_from: Static[str | pathlib.Path | rod.Model | None] = dataclasses.field(
-        repr=False, default=None
+        default=None, repr=False
     )
 
-    _number_of_links: Static[int] = dataclasses.field(
-        init=False, repr=False, default=None
+    integrator: Static[jaxsim.integrators.Integrator | None] = dataclasses.field(
+        default=None, repr=False
     )
 
-    _number_of_joints: Static[int] = dataclasses.field(
-        init=False, repr=False, default=None
+    _description: Static[wrappers.HashlessObject[ModelDescription | None]] = (
+        dataclasses.field(default=None, repr=False)
     )
 
-    def __post_init__(self):
+    @property
+    def description(self) -> ModelDescription:
+        """
+        Return the model description.
+        """
+        return self._description.get()
+
+    def __eq__(self, other: JaxSimModel) -> bool:
+
+        if not isinstance(other, JaxSimModel):
+            return False
+
+        if self.model_name != other.model_name:
+            return False
+
+        if self.time_step != other.time_step:
+            return False
 
-        # These attributes are Static so that we can use `jax.vmap` and `jax.lax.scan`
-        # over the all links and joints
-        with self.mutable_context(
-            mutability=Mutability.MUTABLE_NO_VALIDATION,
-            restore_after_exception=False,
-        ):
-            self._number_of_links = len(self.physics_model.description.links_dict)
-            self._number_of_joints = len(self.physics_model.description.joints_dict)
+        if self.kin_dyn_parameters != other.kin_dyn_parameters:
+            return False
+
+        return True
+
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(self.model_name),
+                hash(float(self.time_step)),
+                hash(self.kin_dyn_parameters),
+                hash(self.contact_model),
+            )
+        )
 
     # ========================
     # Initialization and state
     # ========================
 
-    @staticmethod
+    @classmethod
     def build_from_model_description(
+        cls,
         model_description: str | pathlib.Path | rod.Model,
+        *,
         model_name: str | None = None,
-        gravity: jtp.Array = jaxsim.physics.default_gravity(),
+        time_step: jtp.FloatLike | None = None,
+        integrator: (
+            jaxsim.integrators.Integrator | type[jaxsim.integrators.Integrator] | None
+        ) = None,
+        terrain: jaxsim.terrain.Terrain | None = None,
+        contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         is_urdf: bool | None = None,
-        considered_joints: list[str] | None = None,
+        considered_joints: Sequence[str] | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from a model description.
@@ -81,12 +122,21 @@ class JaxSimModel(JaxsimDataclass):
                 A path to an SDF/URDF file, a string containing
                 its content, or a pre-parsed/pre-built rod model.
             model_name:
-                The optional name of the model that overrides the one in
-                the description.
-            gravity: The 3D gravity vector.
+                The name of the model. If not specified, it is read from the description.
+            time_step:
+                The default time step to consider for the simulation. It can be
+                manually overridden in the function that steps the simulation.
+            terrain: The terrain to consider (the default is a flat infinite plane).
+            contact_model:
+                The contact model to consider.
+                If not specified, a soft contacts model is used.
+            integrator:
+                The integrator to use. If not specified, a default one is used.
+                This argument can either be a pre-built integrator instance or one
+                of the integrator classes defined in JaxSim.
             is_urdf:
-                Whether the model description is a URDF or an SDF. This is
-                automatically inferred if the model description is a path to a file.
+                The optional flag to force the model description to be parsed as a URDF.
+                This is usually automatically inferred.
             considered_joints:
                 The list of joints to consider. If None, all joints are considered.
 
@@ -97,7 +147,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
         )
@@ -109,44 +159,134 @@ class JaxSimModel(JaxsimDataclass):
                 considered_joints=considered_joints
             )
 
-        # Create the physics model from the model description
-        physics_model = jaxsim.physics.model.physics_model.PhysicsModel.build_from(
-            model_description=intermediate_description, gravity=gravity
+        # Build the model.
+        model = cls.build(
+            model_description=intermediate_description,
+            model_name=model_name,
+            time_step=time_step,
+            integrator=integrator,
+            terrain=terrain,
+            contact_model=contact_model,
         )
 
-        # Build the model
-        model = JaxSimModel.build(physics_model=physics_model, model_name=model_name)
-
-        # 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
 
         return model
 
-    @staticmethod
+    @classmethod
     def build(
-        physics_model: jaxsim.physics.model.physics_model.PhysicsModel,
+        cls,
+        model_description: ModelDescription,
+        *,
         model_name: str | None = None,
+        time_step: jtp.FloatLike | None = None,
+        integrator: (
+            jaxsim.integrators.Integrator | type[jaxsim.integrators.Integrator] | None
+        ) = None,
+        terrain: jaxsim.terrain.Terrain | None = None,
+        contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
     ) -> JaxSimModel:
         """
-        Build a Model object from a physics model.
+        Build a Model object from an intermediate model description.
 
         Args:
-            physics_model: The physics model.
+            model_description:
+                The intermediate model description defining the kinematics and dynamics
+                of the model.
             model_name:
+                The name of the model. If not specified, it is read from the description.
+            time_step:
+                The default time step to consider for the simulation. It can be
+                manually overridden in the function that steps the simulation.
+            terrain: The terrain to consider (the default is a flat infinite plane).
                 The optional name of the model overriding the physics model name.
+            integrator:
+                The integrator to use. If not specified, a default one is used.
+                This argument can either be a pre-built integrator instance or one
+                of the integrator classes defined in JaxSim.
+            contact_model:
+                The contact model to consider.
+                If not specified, a soft contacts model is used.
 
         Returns:
             The built Model object.
         """
 
-        # 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 physics_model.description.name
+        # 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
+
+        # Consider the default terrain (a flat infinite plane) if not specified.
+        terrain = (
+            terrain
+            if terrain is not None
+            else JaxSimModel.__dataclass_fields__["terrain"].default_factory()
+        )
+
+        # Consider the default time step if not specified.
+        time_step = (
+            time_step
+            if time_step is not None
+            else JaxSimModel.__dataclass_fields__["time_step"].default_factory()
+        )
+
+        # Create the default contact model.
+        # It will be populated with an initial estimation of good parameters.
+        # While these might not be the best, they are a good starting point.
+        contact_model = (
+            contact_model
+            if contact_model is not None
+            else jaxsim.rbda.contacts.SoftContacts.build()
         )
 
-        # Build the model
-        model = JaxSimModel(physics_model=physics_model, model_name=model_name)  # noqa
+        # Build the integrator if not provided.
+        match integrator:
+
+            # If None, build a default integrator.
+            case None:
+
+                integrator = jaxsim.integrators.fixed_step.Heun2SO3.build(
+                    dynamics=js.ode.wrap_system_dynamics_for_integration(
+                        system_dynamics=js.ode.system_dynamics
+                    )
+                )
+
+            # If it's a pre-built integrator (also a custom one from the user)
+            # just use it as is.
+            case _ if isinstance(integrator, jaxsim.integrators.Integrator):
+                pass
+
+            # If an integrator class is passed, assume that it is a JaxSim integrator
+            # and build it with the default system dynamics.
+            case _ if issubclass(integrator, jaxsim.integrators.Integrator):
+
+                integrator_cls = integrator
+                integrator = integrator_cls.build(
+                    dynamics=js.ode.wrap_system_dynamics_for_integration(
+                        system_dynamics=js.ode.system_dynamics
+                    )
+                )
+
+            case _:
+                raise ValueError(f"Invalid integrator: {integrator}")
+
+        # Build the model.
+        model = cls(
+            model_name=model_name,
+            kin_dyn_parameters=js.kin_dyn_parameters.KinDynParameters.build(
+                model_description=model_description
+            ),
+            time_step=time_step,
+            terrain=terrain,
+            contact_model=contact_model,
+            integrator=integrator,
+            # The following is wrapped as hashless since it's a static argument, and we
+            # don't want to trigger recompilation if it changes. All relevant parameters
+            # needed to compute kinematics and dynamics quantities are stored in the
+            # kin_dyn_parameters attribute.
+            _description=wrappers.HashlessObject(obj=model_description),
+        )
 
         return model
 
@@ -164,7 +304,7 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.model_name
 
-    def number_of_links(self) -> jtp.Int:
+    def number_of_links(self) -> int:
         """
         Return the number of links in the model.
 
@@ -175,9 +315,9 @@ class JaxSimModel(JaxsimDataclass):
             The base link is included in the count and its index is always 0.
         """
 
-        return self._number_of_links
+        return self.kin_dyn_parameters.number_of_links()
 
-    def number_of_joints(self) -> jtp.Int:
+    def number_of_joints(self) -> int:
         """
         Return the number of joints in the model.
 
@@ -185,7 +325,18 @@ class JaxSimModel(JaxsimDataclass):
             The number of joints in the model.
         """
 
-        return self._number_of_joints
+        return self.kin_dyn_parameters.number_of_joints()
+
+    def number_of_frames(self) -> int:
+        """
+        Return the number of frames in the model.
+
+        Returns:
+            The number of frames in the model.
+
+        """
+
+        return self.kin_dyn_parameters.number_of_frames()
 
     # =================
     # Base link methods
@@ -199,7 +350,7 @@ class JaxSimModel(JaxsimDataclass):
             True if the model is floating-base, False otherwise.
         """
 
-        return self.physics_model.is_floating_base
+        return bool(self.kin_dyn_parameters.joint_model.joint_dofs[0] == 6)
 
     def base_link(self) -> str:
         """
@@ -207,9 +358,12 @@ class JaxSimModel(JaxsimDataclass):
 
         Returns:
             The name of the base link.
+
+        Note:
+            By default, the base link is the root of the kinematic tree.
         """
 
-        return self.physics_model.description.root.name
+        return self.link_names()[0]
 
     # =====================
     # Joint-related methods
@@ -227,7 +381,7 @@ class JaxSimModel(JaxsimDataclass):
             the number of joints. In the future, this could be different.
         """
 
-        return len(self.physics_model.description.joints_dict)
+        return int(sum(self.kin_dyn_parameters.joint_model.joint_dofs[1:]))
 
     def joint_names(self) -> tuple[str, ...]:
         """
@@ -237,7 +391,7 @@ class JaxSimModel(JaxsimDataclass):
             The names of the joints in the model.
         """
 
-        return tuple(self.physics_model.description.joints_dict.keys())
+        return self.kin_dyn_parameters.joint_model.joint_names[1:]
 
     # ====================
     # Link-related methods
@@ -251,7 +405,21 @@ class JaxSimModel(JaxsimDataclass):
             The names of the links in the model.
         """
 
-        return tuple(self.physics_model.description.links_dict.keys())
+        return self.kin_dyn_parameters.link_names
+
+    # =====================
+    # Frame-related methods
+    # =====================
+
+    def frame_names(self) -> tuple[str, ...]:
+        """
+        Return the names of the frames in the model.
+
+        Returns:
+            The names of the frames in the model.
+        """
+
+        return self.kin_dyn_parameters.frame_parameters.name
 
 
 # =====================
@@ -259,42 +427,63 @@ class JaxSimModel(JaxsimDataclass):
 # =====================
 
 
-def reduce(model: JaxSimModel, considered_joints: tuple[str, ...]) -> JaxSimModel:
+def reduce(
+    model: JaxSimModel,
+    considered_joints: tuple[str, ...],
+    locked_joint_positions: dict[str, jtp.FloatLike] | None = None,
+) -> JaxSimModel:
     """
     Reduce the model by lumping together the links connected by removed joints.
 
     Args:
         model: The model to reduce.
         considered_joints: The sequence of joints to consider.
-
-    Note:
-        If considered_joints contains joints not existing in the model, the method
-        will raise an exception. If considered_joints is empty, the method will
-        return a copy of the input model.
+        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 of this dictionary.
+            If a joint is not part of the dictionary, its position is set to zero.
     """
 
-    if len(considered_joints) == 0:
-        return model.copy()
+    locked_joint_positions = (
+        locked_joint_positions if locked_joint_positions is not None else {}
+    )
+
+    # If locked joints are passed, make sure that they are valid.
+    if not set(locked_joint_positions).issubset(model.joint_names()):
+        new_joints = set(model.joint_names()) - set(locked_joint_positions)
+        raise ValueError(f"Passed joints not existing in the model: {new_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
+    # to removed joints.
+    for joint_name in set(model.joint_names()) - set(considered_joints):
+        j = intermediate_description.joints_dict[joint_name]
+        with j.mutable_context():
+            j.initial_position = float(locked_joint_positions.get(joint_name, 0.0))
 
     # Reduce the model description.
-    # If considered_joints contains joints not existing in the model, the method
-    # will raise an exception.
-    reduced_intermediate_description = model.physics_model.description.reduce(
+    # If `considered_joints` contains joints not existing in the model,
+    # the method will raise an exception.
+    reduced_intermediate_description = intermediate_description.reduce(
         considered_joints=list(considered_joints)
     )
 
-    # Create the physics model from the reduced model description
-    physics_model = jaxsim.physics.model.physics_model.PhysicsModel.build_from(
-        model_description=reduced_intermediate_description,
-        gravity=model.physics_model.gravity[0:3],
-    )
-
-    # Build the reduced model
+    # Build the reduced model.
     reduced_model = JaxSimModel.build(
-        physics_model=physics_model, model_name=model.name()
+        model_description=reduced_intermediate_description,
+        model_name=model.name(),
+        time_step=model.time_step,
+        terrain=model.terrain,
+        contact_model=model.contact_model,
+        integrator=model.integrator,
     )
 
-    # 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
 
@@ -307,6 +496,7 @@ def reduce(model: JaxSimModel, considered_joints: tuple[str, ...]) -> JaxSimMode
 
 
 @jax.jit
+@js.common.named_scope
 def total_mass(model: JaxSimModel) -> jtp.Float:
     """
     Compute the total mass of the model.
@@ -318,52 +508,25 @@ def total_mass(model: JaxSimModel) -> jtp.Float:
         The total mass of the model.
     """
 
-    return (
-        jax.vmap(lambda idx: js.link.mass(model=model, link_index=idx))(
-            jnp.arange(model.number_of_links())
-        )
-        .sum()
-        .astype(float)
-    )
-
-
-# ==============
-# Center of mass
-# ==============
+    return model.kin_dyn_parameters.link_parameters.mass.sum().astype(float)
 
 
 @jax.jit
-def com_position(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vector:
+@js.common.named_scope
+def link_spatial_inertia_matrices(model: JaxSimModel) -> jtp.Array:
     """
-    Compute the position of the center of mass of the model.
+    Compute the spatial 6D inertia matrices of all links of the model.
 
     Args:
         model: The model to consider.
-        data: The data of the considered model.
 
     Returns:
-        The position of the center of mass of the model w.r.t. the world frame.
+        A 3D array containing the stacked spatial 6D inertia matrices of the links.
     """
 
-    m = total_mass(model=model)
-
-    W_H_L = forward_kinematics(model=model, data=data)
-    W_H_B = data.base_transform()
-    B_H_W = jaxlie.SE3.from_matrix(W_H_B).inverse().as_matrix()
-
-    def B_p̃_LCoM(i) -> jtp.Vector:
-        m = js.link.mass(model=model, link_index=i)
-        L_p_LCoM = js.link.com_position(
-            model=model, data=data, link_index=i, in_link_frame=True
-        )
-        return m * B_H_W @ W_H_L[i] @ jnp.hstack([L_p_LCoM, 1])
-
-    com_links = jax.vmap(B_p̃_LCoM)(jnp.arange(model.number_of_links()))
-
-    B_p̃_CoM = (1 / m) * com_links.sum(axis=0)
-    B_p̃_CoM = B_p̃_CoM.at[3].set(1)
-
-    return (W_H_B @ B_p̃_CoM)[0:3].astype(float)
+    return jax.vmap(js.kin_dyn_parameters.LinkParameters.spatial_inertia)(
+        model.kin_dyn_parameters.link_parameters
+    )
 
 
 # ==============================
@@ -372,6 +535,7 @@ def com_position(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vecto
 
 
 @jax.jit
+@js.common.named_scope
 def forward_kinematics(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
     """
     Compute the SE(3) transforms from the world frame to the frames of all links.
@@ -385,10 +549,11 @@ def forward_kinematics(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp
         The first axis is the link index.
     """
 
-    W_H_LL = jaxsim.physics.algos.forward_kinematics.forward_kinematics_model(
-        model=model.physics_model,
-        q=data.state.physics_model.joint_positions,
-        xfb=data.state.physics_model.xfb(),
+    W_H_LL = jaxsim.rbda.forward_kinematics_model(
+        model=model,
+        base_position=data.base_position(),
+        base_quaternion=data.base_orientation(dcm=False),
+        joint_positions=data.joint_positions(model=model),
     )
 
     return jnp.atleast_3d(W_H_LL).astype(float)
@@ -424,51 +589,296 @@ def generalized_free_floating_jacobian(
         output_vel_repr if output_vel_repr is not None else data.velocity_representation
     )
 
-    # The body frame of the link.jacobian method is the link frame L.
-    # In this method, we want instead to use the base link B as body frame.
-    # Therefore, we always get the link jacobian having Inertial as output
-    # representation, and then we convert it to the desired output representation.
+    # Compute the doubly-left free-floating full jacobian.
+    B_J_full_WX_B, B_H_L = jaxsim.rbda.jacobian_full_doubly_left(
+        model=model,
+        joint_positions=data.joint_positions(),
+    )
+
+    # ======================================================================
+    # Update the input velocity representation such that v_WL = J_WL_I @ I_ν
+    # ======================================================================
+
+    match data.velocity_representation:
+
+        case VelRepr.Inertial:
+
+            W_H_B = data.base_transform()
+            B_X_W = Adjoint.from_transform(transform=W_H_B, inverse=True)
+
+            B_J_full_WX_I = B_J_full_WX_W = (  # noqa: F841
+                B_J_full_WX_B
+                @ jax.scipy.linalg.block_diag(B_X_W, jnp.eye(model.dofs()))
+            )
+
+        case VelRepr.Body:
+
+            B_J_full_WX_I = B_J_full_WX_B
+
+        case VelRepr.Mixed:
+
+            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 = Adjoint.from_transform(transform=BW_H_B, inverse=True)
+
+            B_J_full_WX_I = B_J_full_WX_BW = (  # noqa: F841
+                B_J_full_WX_B
+                @ jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
+            )
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # ====================================================================
+    # Create stacked Jacobian for each link by filtering the full Jacobian
+    # ====================================================================
+
+    κ_bool = model.kin_dyn_parameters.support_body_array_bool
+
+    # Keep only the columns of the full Jacobian corresponding to the support
+    # body array of each link.
+    B_J_WL_I = jax.vmap(
+        lambda κ: jnp.where(
+            jnp.hstack([jnp.ones(5), κ]), B_J_full_WX_I, jnp.zeros_like(B_J_full_WX_I)
+        )
+    )(κ_bool)
+
+    # =======================================================================
+    # Update the output velocity representation such that O_v_WL = O_J_WL @ ν
+    # =======================================================================
+
     match output_vel_repr:
+
         case VelRepr.Inertial:
-            to_output = lambda W_J_WL: W_J_WL
+
+            W_H_B = data.base_transform()
+            W_X_B = jaxsim.math.Adjoint.from_transform(W_H_B)
+
+            O_J_WL_I = W_J_WL_I = jax.vmap(  # noqa: F841
+                lambda B_J_WL_I: W_X_B @ B_J_WL_I
+            )(B_J_WL_I)
 
         case VelRepr.Body:
 
-            def to_output(W_J_WL: jtp.Matrix) -> jtp.Matrix:
-                W_H_B = data.base_transform()
-                B_X_W = jaxlie.SE3.from_matrix(W_H_B).inverse().adjoint()
-                return B_X_W @ W_J_WL
+            O_J_WL_I = L_J_WL_I = jax.vmap(  # noqa: F841
+                lambda B_H_L, B_J_WL_I: jaxsim.math.Adjoint.from_transform(
+                    B_H_L, inverse=True
+                )
+                @ B_J_WL_I
+            )(B_H_L, B_J_WL_I)
 
         case VelRepr.Mixed:
 
-            def to_output(W_J_WL: jtp.Matrix) -> jtp.Matrix:
-                W_H_B = data.base_transform()
-                W_H_BW = jnp.array(W_H_B).at[0:3, 0:3].set(jnp.eye(3))
-                BW_X_W = jaxlie.SE3.from_matrix(W_H_BW).inverse().adjoint()
-                return BW_X_W @ W_J_WL
+            W_H_B = data.base_transform()
+
+            LW_H_L = jax.vmap(
+                lambda B_H_L: (W_H_B @ B_H_L).at[0:3, 3].set(jnp.zeros(3))
+            )(B_H_L)
+
+            LW_H_B = jax.vmap(
+                lambda LW_H_L, B_H_L: LW_H_L @ jaxsim.math.Transform.inverse(B_H_L)
+            )(LW_H_L, B_H_L)
+
+            O_J_WL_I = LW_J_WL_I = jax.vmap(  # noqa: F841
+                lambda LW_H_B, B_J_WL_I: jaxsim.math.Adjoint.from_transform(LW_H_B)
+                @ B_J_WL_I
+            )(LW_H_B, B_J_WL_I)
 
         case _:
             raise ValueError(output_vel_repr)
 
-    # Compute first the link jacobians having the active representation of `data`
-    # as input representation (matching the one of ν), and inertial as output
-    # representation (i.e. W_J_WL_C where C is C_ν).
-    # Then, with to_output, we convert this jacobian to the desired output
-    # representation, that can either be W (inertial), B (body), or B[W] (mixed).
-    # This is necessary because for example the body-fixed free-floating jacobian
-    # of a link is L_J_WL, but here being inside model we need B_J_WL.
-    J_free_floating = jax.vmap(
-        lambda i: to_output(
-            W_J_WL=js.link.jacobian(
-                model=model,
-                data=data,
-                link_index=i,
-                output_vel_repr=VelRepr.Inertial,
+    return O_J_WL_I
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def generalized_free_floating_jacobian_derivative(
+    model: JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    """
+    Compute the free-floating jacobian derivatives of all links.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        output_vel_repr:
+            The output velocity representation of the free-floating jacobian derivatives.
+
+    Returns:
+        The `(nL, 6, 6+dofs)` array containing the stacked free-floating
+        jacobian derivatives of the links. The first axis is the link index.
+    """
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Compute the derivative of the doubly-left free-floating full jacobian.
+    B_J̇_full_WX_B, B_H_L = jaxsim.rbda.jacobian_derivative_full_doubly_left(
+        model=model,
+        joint_positions=data.joint_positions(),
+        joint_velocities=data.joint_velocities(),
+    )
+
+    # The derivative of the equation to change the input and output representations
+    # of the Jacobian derivative needs the computation of the plain link Jacobian.
+    B_J_full_WL_B, _ = jaxsim.rbda.jacobian_full_doubly_left(
+        model=model,
+        joint_positions=data.joint_positions(),
+    )
+
+    # Compute the actual doubly-left free-floating jacobian derivative of the link
+    # by zeroing the columns not in the path π_B(L) using the boolean κ(i).
+    κb = model.kin_dyn_parameters.support_body_array_bool
+
+    # Compute the base transform.
+    W_H_B = data.base_transform()
+
+    # We add the 5 columns of ones to the Jacobian derivative to account for the
+    # base velocity and acceleration (5 + number of links = 6 + number of joints).
+    B_J̇_WL_B = (
+        jnp.hstack([jnp.ones((κb.shape[0], 5)), κb])[:, jnp.newaxis] * B_J̇_full_WX_B
+    )
+    B_J_WL_B = (
+        jnp.hstack([jnp.ones((κb.shape[0], 5)), κb])[:, jnp.newaxis] * B_J_full_WL_B
+    )
+
+    # =====================================================
+    # Compute quantities to adjust the input representation
+    # =====================================================
+
+    In = jnp.eye(model.dofs())
+    On = jnp.zeros(shape=(model.dofs(), model.dofs()))
+
+    match data.velocity_representation:
+
+        case VelRepr.Inertial:
+
+            B_X_W = jaxsim.math.Adjoint.from_transform(transform=W_H_B, inverse=True)
+
+            W_v_WB = data.base_velocity()
+            B_Ẋ_W = -B_X_W @ jaxsim.math.Cross.vx(W_v_WB)
+
+            # Compute the operator to change the representation of ν, and its
+            # time derivative.
+            T = jax.scipy.linalg.block_diag(B_X_W, In)
+            Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_W, On)
+
+        case VelRepr.Body:
+
+            B_X_B = jaxsim.math.Adjoint.from_rotation_and_translation(
+                translation=jnp.zeros(3), rotation=jnp.eye(3)
+            )
+
+            B_Ẋ_B = jnp.zeros(shape=(6, 6))
+
+            # Compute the operator to change the representation of ν, and its
+            # time derivative.
+            T = jax.scipy.linalg.block_diag(B_X_B, In)
+            Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_B, On)
+
+        case VelRepr.Mixed:
+
+            BW_H_B = W_H_B.at[0:3, 3].set(jnp.zeros(3))
+            B_X_BW = jaxsim.math.Adjoint.from_transform(transform=BW_H_B, inverse=True)
+
+            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)
+
+            # Compute the operator to change the representation of ν, and its
+            # time derivative.
+            T = jax.scipy.linalg.block_diag(B_X_BW, In)
+            Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_BW, On)
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # ======================================================
+    # Compute quantities to adjust the output representation
+    # ======================================================
+
+    match output_vel_repr:
+
+        case VelRepr.Inertial:
+
+            O_X_B = W_X_B = jaxsim.math.Adjoint.from_transform(transform=W_H_B)
+
+            with data.switch_velocity_representation(VelRepr.Body):
+                B_v_WB = data.base_velocity()
+
+            O_Ẋ_B = W_Ẋ_B = W_X_B @ jaxsim.math.Cross.vx(B_v_WB)  # noqa: F841
+
+        case VelRepr.Body:
+
+            O_X_B = L_X_B = jaxsim.math.Adjoint.from_transform(
+                transform=B_H_L, inverse=True
+            )
+
+            B_X_L = jaxsim.math.Adjoint.inverse(adjoint=L_X_B)
+
+            with data.switch_velocity_representation(VelRepr.Body):
+                B_v_WB = data.base_velocity()
+                L_v_WL = jnp.einsum(
+                    "b6j,j->b6", L_X_B @ B_J_WL_B, data.generalized_velocity()
+                )
+
+            O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
+                jnp.einsum("bij,bj->bi", B_X_L, L_v_WL) - B_v_WB
+            )
+
+        case VelRepr.Mixed:
+
+            W_H_L = W_H_B @ B_H_L
+            LW_H_L = W_H_L.at[:, 0:3, 3].set(jnp.zeros_like(W_H_L[:, 0:3, 3]))
+            LW_H_B = LW_H_L @ jaxsim.math.Transform.inverse(B_H_L)
+
+            O_X_B = LW_X_B = jaxsim.math.Adjoint.from_transform(transform=LW_H_B)
+
+            B_X_LW = jaxsim.math.Adjoint.inverse(adjoint=LW_X_B)
+
+            with data.switch_velocity_representation(VelRepr.Body):
+                B_v_WB = data.base_velocity()
+
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                BW_H_B = W_H_B.at[0:3, 3].set(jnp.zeros(3))
+                B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
+                LW_v_WL = jnp.einsum(
+                    "bij,bj->bi",
+                    LW_X_B,
+                    B_J_WL_B
+                    @ jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
+                    @ data.generalized_velocity(),
+                )
+
+                LW_v_W_LW = LW_v_WL.at[:, 3:6].set(jnp.zeros_like(LW_v_WL[:, 3:6]))
+
+            LW_v_LW_L = LW_v_WL - LW_v_W_LW
+            LW_v_B_LW = LW_v_WL - jnp.einsum("bij,j->bi", LW_X_B, B_v_WB) - LW_v_LW_L
+
+            O_Ẋ_B = LW_Ẋ_B = -LW_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
+                jnp.einsum("bij,bj->bi", B_X_LW, LW_v_B_LW)
             )
-        )
-    )(jnp.arange(model.number_of_links()))
 
-    return J_free_floating
+        case _:
+            raise ValueError(output_vel_repr)
+
+    # =============================================================
+    # Express the Jacobian derivative in the target representations
+    # =============================================================
+
+    # Sum all the components that form the Jacobian derivative in the target
+    # input/output velocity representations.
+    O_J̇_WL_I = jnp.zeros_like(B_J̇_WL_B)
+    O_J̇_WL_I += O_Ẋ_B @ B_J_WL_B @ T
+    O_J̇_WL_I += O_X_B @ B_J̇_WL_B @ T
+    O_J̇_WL_I += O_X_B @ B_J_WL_B @ Ṫ
+
+    return O_J̇_WL_I
 
 
 @functools.partial(jax.jit, static_argnames=["prefer_aba"])
@@ -477,7 +887,7 @@ def forward_dynamics(
     data: js.data.JaxSimModelData,
     *,
     joint_forces: jtp.VectorLike | None = None,
-    external_forces: jtp.MatrixLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
     prefer_aba: float = True,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
@@ -488,8 +898,8 @@ def forward_dynamics(
         data: The data of the considered model.
         joint_forces:
             The joint forces to consider as a vector of shape `(dofs,)`.
-        external_forces:
-            The external forces to consider as a matrix of shape `(nL, 6)`.
+        link_forces:
+            The link 6D forces consider as a matrix of shape `(nL, 6)`.
             The frame in which they are expressed must be `data.velocity_representation`.
         prefer_aba: Whether to prefer the ABA algorithm over the CRB one.
 
@@ -505,17 +915,18 @@ def forward_dynamics(
         model=model,
         data=data,
         joint_forces=joint_forces,
-        external_forces=external_forces,
+        link_forces=link_forces,
     )
 
 
 @jax.jit
+@js.common.named_scope
 def forward_dynamics_aba(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
     joint_forces: jtp.VectorLike | None = None,
-    external_forces: jtp.MatrixLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
     Compute the forward dynamics of the model with the ABA algorithm.
@@ -525,8 +936,8 @@ def forward_dynamics_aba(
         data: The data of the considered model.
         joint_forces:
             The joint forces to consider as a vector of shape `(dofs,)`.
-        external_forces:
-            The external forces to consider as a matrix of shape `(nL, 6)`.
+        link_forces:
+            The link 6D forces to consider as a matrix of shape `(nL, 6)`.
             The frame in which they are expressed must be `data.velocity_representation`.
 
     Returns:
@@ -535,86 +946,132 @@ def forward_dynamics_aba(
         considered joint forces and external forces.
     """
 
-    # Build joint torques if not provided
+    # ============
+    # Prepare data
+    # ============
+
+    # Build joint forces, if not provided.
     τ = (
-        joint_forces
+        jnp.atleast_1d(joint_forces.squeeze())
         if joint_forces is not None
         else jnp.zeros_like(data.joint_positions())
     )
 
-    # Build external forces if not provided
-    f_ext = (
-        external_forces
-        if external_forces is not None
+    # Build link forces, if not provided.
+    f_L = (
+        jnp.atleast_2d(link_forces.squeeze())
+        if link_forces is not None
         else jnp.zeros((model.number_of_links(), 6))
     )
 
-    # Compute ABA
-    W_v̇_WB, s̈ = jaxsim.physics.algos.aba.aba(
-        model=model.physics_model,
-        xfb=data.state.physics_model.xfb(),
-        q=data.state.physics_model.joint_positions,
-        qd=data.state.physics_model.joint_velocities,
-        tau=τ,
-        f_ext=f_ext,
+    # Create a references object that simplifies converting among representations.
+    references = js.references.JaxSimModelReferences.build(
+        model=model,
+        joint_force_references=τ,
+        link_forces=f_L,
+        data=data,
+        velocity_representation=data.velocity_representation,
     )
 
-    def to_active(W_vd_WB, W_H_C, W_v_WB, W_vl_WC):
-        C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
+    # Extract the link and joint serializations.
+    link_names = model.link_names()
+    joint_names = model.joint_names()
+
+    # Extract the state in inertial-fixed representation.
+    with data.switch_velocity_representation(VelRepr.Inertial):
+        W_p_B = data.base_position()
+        W_v_WB = data.base_velocity()
+        W_Q_B = data.base_orientation(dcm=False)
+        s = data.joint_positions(model=model, joint_names=joint_names)
+        ṡ = data.joint_velocities(model=model, joint_names=joint_names)
+
+    # Extract the inputs in inertial-fixed representation.
+    with references.switch_velocity_representation(VelRepr.Inertial):
+        W_f_L = references.link_forces(model=model, data=data, link_names=link_names)
+        τ = references.joint_force_references(model=model, joint_names=joint_names)
+
+    # ========================
+    # Compute forward dynamics
+    # ========================
+
+    W_v̇_WB, s̈ = jaxsim.rbda.aba(
+        model=model,
+        base_position=W_p_B,
+        base_quaternion=W_Q_B,
+        joint_positions=s,
+        base_linear_velocity=W_v_WB[0:3],
+        base_angular_velocity=W_v_WB[3:6],
+        joint_velocities=ṡ,
+        joint_forces=τ,
+        link_forces=W_f_L,
+        standard_gravity=data.standard_gravity(),
+    )
 
-        if data.velocity_representation != VelRepr.Mixed:
-            return C_X_W @ W_vd_WB
+    # =============
+    # Adjust output
+    # =============
 
-        from jaxsim.math.cross import Cross
+    def to_active(
+        W_v̇_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WB: jtp.Vector, W_v_WC: jtp.Vector
+    ) -> jtp.Vector:
+        """
+        Convert the inertial-fixed apparent base acceleration W_v̇_WB to
+        another representation C_v̇_WB expressed in a generic frame C.
+        """
 
-        W_v_WC = jnp.hstack([W_vl_WC, jnp.zeros(3)])
-        return C_X_W @ (W_vd_WB - Cross.vx(W_v_WC) @ W_v_WB)
+        # 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 = 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:
         case VelRepr.Inertial:
-            W_H_C = W_H_W = jnp.eye(4)
-            W_vl_WC = W_vl_WW = jnp.zeros(3)
+            # In this case C=W
+            W_H_C = W_H_W = jnp.eye(4)  # noqa: F841
+            W_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
 
         case VelRepr.Body:
+            # In this case C=B
             W_H_C = W_H_B = data.base_transform()
-            W_vl_WC = W_vl_WB = data.base_velocity()[0:3]
+            W_v_WC = W_v_WB
 
         case VelRepr.Mixed:
+            # In this case C=B[W]
             W_H_B = data.base_transform()
-            W_H_C = W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
-            W_vl_WC = W_vl_W_BW = data.base_velocity()[0:3]
+            W_H_C = W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))  # noqa: F841
+            W_ṗ_B = data.base_velocity()[0:3]
+            W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)  # noqa: F841
 
         case _:
             raise ValueError(data.velocity_representation)
 
-    # We need to convert the derivative of the base acceleration to the active
+    # We need to convert the derivative of the base velocity to the active
     # representation. In Mixed representation, this conversion is not a plain
     # transformation with just X, but it also involves a cross product in ℝ⁶.
     C_v̇_WB = to_active(
-        W_vd_WB=W_v̇_WB.squeeze(),
+        W_v̇_WB=W_v̇_WB,
         W_H_C=W_H_C,
-        W_v_WB=jnp.hstack(
-            [
-                data.state.physics_model.base_linear_velocity,
-                data.state.physics_model.base_angular_velocity,
-            ]
-        ),
-        W_vl_WC=W_vl_WC,
+        W_v_WB=W_v_WB,
+        W_v_WC=W_v_WC,
     )
 
-    # Adjust shape
-    s̈ = jnp.atleast_1d(s̈.squeeze())
+    # The ABA algorithm already returns a zero base 6D acceleration for
+    # fixed-based models. However, the to_active function introduces an
+    # additional acceleration component in Mixed representation.
+    # Here below we make sure that the base acceleration is zero.
+    C_v̇_WB = C_v̇_WB if model.floating_base() else jnp.zeros(6)
 
-    return C_v̇_WB, s̈
+    return C_v̇_WB.astype(float), s̈.astype(float)
 
 
 @jax.jit
+@js.common.named_scope
 def forward_dynamics_crb(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
     joint_forces: jtp.VectorLike | None = None,
-    external_forces: jtp.MatrixLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
     Compute the forward dynamics of the model with the CRB algorithm.
@@ -624,8 +1081,8 @@ def forward_dynamics_crb(
         data: The data of the considered model.
         joint_forces:
             The joint forces to consider as a vector of shape `(dofs,)`.
-        external_forces:
-            The external forces to consider as a matrix of shape `(nL, 6)`.
+        link_forces:
+            The link 6D forces to consider as a matrix of shape `(nL, 6)`.
             The frame in which they are expressed must be `data.velocity_representation`.
 
     Returns:
@@ -638,21 +1095,25 @@ def forward_dynamics_crb(
         models with a large number of degrees of freedom.
     """
 
-    # Build joint torques if not provided
+    # ============
+    # Prepare data
+    # ============
+
+    # 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(external_forces)
-        if external_forces is not None
+        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
@@ -660,6 +1121,10 @@ def forward_dynamics_crb(
 
     # TODO: invert the Mss block exploiting sparsity defined by the parent array λ(i)
 
+    # ========================
+    # Compute forward dynamics
+    # ========================
+
     if model.floating_base():
         # l: number of links.
         # g: generalized coordinates, 6 + number of joints.
@@ -675,19 +1140,24 @@ def forward_dynamics_crb(
         v̇_WB = jnp.zeros(6)
         ν̇ = jnp.hstack([v̇_WB, s̈.squeeze()])
 
+    # =============
+    # Adjust output
+    # =============
+
     # Extract the base acceleration in the active representation.
     # Note that this is an apparent acceleration (relevant in Mixed representation),
     # therefore it cannot be always expressed in different frames with just a
     # 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̈
 
 
 @jax.jit
+@js.common.named_scope
 def free_floating_mass_matrix(
     model: JaxSimModel, data: js.data.JaxSimModelData
 ) -> jtp.Matrix:
@@ -702,9 +1172,9 @@ def free_floating_mass_matrix(
         The free-floating mass matrix of the model.
     """
 
-    M_body = jaxsim.physics.algos.crba.crba(
-        model=model.physics_model,
-        q=data.state.physics_model.joint_positions,
+    M_body = jaxsim.rbda.crba(
+        model=model,
+        joint_positions=data.state.physics_model.joint_positions,
     )
 
     match data.velocity_representation:
@@ -712,29 +1182,19 @@ def free_floating_mass_matrix(
             return M_body
 
         case VelRepr.Inertial:
-            zero_6n = jnp.zeros(shape=(6, model.dofs()))
-            B_X_W = jaxlie.SE3.from_matrix(data.base_transform()).inverse().adjoint()
 
-            invT = jnp.vstack(
-                [
-                    jnp.block([B_X_W, zero_6n]),
-                    jnp.block([zero_6n.T, jnp.eye(model.dofs())]),
-                ]
+            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
 
         case VelRepr.Mixed:
-            zero_6n = jnp.zeros(shape=(6, model.dofs()))
-            W_H_BW = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
-            BW_X_W = jaxlie.SE3.from_matrix(W_H_BW).inverse().adjoint()
-
-            invT = jnp.vstack(
-                [
-                    jnp.block([BW_X_W, zero_6n]),
-                    jnp.block([zero_6n.T, jnp.eye(model.dofs())]),
-                ]
-            )
+
+            BW_H_B = data.base_transform().at[0:3, 3].set(jnp.zeros(3))
+            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
 
@@ -743,77 +1203,206 @@ def free_floating_mass_matrix(
 
 
 @jax.jit
-def inverse_dynamics(
-    model: JaxSimModel,
-    data: js.data.JaxSimModelData,
-    *,
-    joint_accelerations: jtp.Vector | None = None,
-    base_acceleration: jtp.Vector | None = None,
-    external_forces: jtp.Matrix | None = None,
-) -> tuple[jtp.Vector, jtp.Vector]:
+@js.common.named_scope
+def free_floating_coriolis_matrix(
+    model: JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
     """
-    Compute inverse dynamics with the RNEA algorithm.
+    Compute the free-floating Coriolis matrix of the model.
 
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_accelerations:
-            The joint accelerations to consider as a vector of shape `(dofs,)`.
-        base_acceleration:
-            The base acceleration to consider as a vector of shape `(6,)`.
-        external_forces:
-            The external forces to consider as a matrix of shape `(nL, 6)`.
-            The frame in which they are expressed must be `data.velocity_representation`.
 
     Returns:
-        A tuple containing the 6D force in the active representation applied to the
-        base to obtain the considered base acceleration, and the joint forces to apply
-        to obtain the considered joint accelerations.
+        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.
     """
 
-    # Build joint accelerations if not provided
-    joint_accelerations = (
-        joint_accelerations
-        if joint_accelerations is not None
-        else jnp.zeros_like(data.joint_positions())
-    )
+    # 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):
 
-    # Build base acceleration if not provided
-    base_acceleration = (
-        base_acceleration if base_acceleration is not None else jnp.zeros(6)
-    )
+        B_ν = data.generalized_velocity()
 
-    external_forces = (
-        external_forces
-        if external_forces is not None
-        else jnp.zeros(shape=(model.number_of_links(), 6))
-    )
+        # Doubly-left free-floating Jacobian.
+        L_J_WL_B = generalized_free_floating_jacobian(model=model, data=data)
 
-    def to_inertial(C_v̇_WB, W_H_C, C_v_WB, W_vl_WC):
-        W_X_C = jaxlie.SE3.from_matrix(W_H_C).adjoint()
-        C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
+        # Doubly-left free-floating Jacobian derivative.
+        L_J̇_WL_B = generalized_free_floating_jacobian_derivative(model=model, data=data)
 
-        if data.velocity_representation != VelRepr.Mixed:
-            return W_X_C @ C_v̇_WB
-        else:
-            from jaxsim.math.cross import Cross
+    L_M_L = link_spatial_inertia_matrices(model=model)
 
-            C_v_WC = C_X_W @ jnp.hstack([W_vl_WC, jnp.zeros(3)])
-            return W_X_C @ (C_v̇_WB + Cross.vx(C_v_WC) @ C_v_WB)
+    # 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)
 
-    match data.velocity_representation:
-        case VelRepr.Inertial:
-            W_H_C = W_H_W = jnp.eye(4)
-            W_vl_WC = W_vl_WW = jnp.zeros(3)
+    # Compute the contribution of each link to the Coriolis matrix.
+    def compute_link_contribution(M, v, J, J̇) -> jtp.Array:
 
-        case VelRepr.Body:
-            W_H_C = W_H_B = data.base_transform()
-            W_vl_WC = W_vl_WB = data.base_velocity()[0:3]
+        return J.T @ ((Cross.vx_star(v) @ M + M @ Cross.vx(v)) @ J + M @ J̇)
 
-        case VelRepr.Mixed:
-            W_H_B = data.base_transform()
-            W_H_C = W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
-            W_vl_WC = W_vl_W_BW = data.base_velocity()[0:3]
+    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
+@js.common.named_scope
+def inverse_dynamics(
+    model: JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    joint_accelerations: jtp.VectorLike | None = None,
+    base_acceleration: jtp.VectorLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
+) -> tuple[jtp.Vector, jtp.Vector]:
+    """
+    Compute inverse dynamics with the RNEA algorithm.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        joint_accelerations:
+            The joint accelerations to consider as a vector of shape `(dofs,)`.
+        base_acceleration:
+            The base acceleration to consider as a vector of shape `(6,)`.
+        link_forces:
+            The link 6D forces to consider as a matrix of shape `(nL, 6)`.
+            The frame in which they are expressed must be `data.velocity_representation`.
+
+    Returns:
+        A tuple containing the 6D force in the active representation applied to the
+        base to obtain the considered base acceleration, and the joint forces to apply
+        to obtain the considered joint accelerations.
+    """
+
+    # ============
+    # Prepare data
+    # ============
+
+    # Build joint accelerations, if not provided.
+    s̈ = (
+        jnp.atleast_1d(jnp.array(joint_accelerations).squeeze())
+        if joint_accelerations is not None
+        else jnp.zeros_like(data.joint_positions())
+    )
+
+    # Build base acceleration, if not provided.
+    v̇_WB = (
+        jnp.array(base_acceleration).squeeze()
+        if base_acceleration is not None
+        else jnp.zeros(6)
+    )
+
+    # Build link forces, if not provided.
+    f_L = (
+        jnp.atleast_2d(jnp.array(link_forces).squeeze())
+        if link_forces is not None
+        else jnp.zeros(shape=(model.number_of_links(), 6))
+    )
+
+    def to_inertial(C_v̇_WB, W_H_C, C_v_WB, W_v_WC):
+        """
+        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.
+        """
+
+        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 ℝ⁶.
+        # In Inertial and Body representations, the cross product is always zero.
+        return W_X_C @ (C_v̇_WB + Cross.vx(C_v_WC) @ C_v_WB)
+
+    match data.velocity_representation:
+        case VelRepr.Inertial:
+            W_H_C = W_H_W = jnp.eye(4)  # noqa: F841
+            W_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
+
+        case VelRepr.Body:
+            W_H_C = W_H_B = data.base_transform()
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                W_v_WC = W_v_WB = data.base_velocity()
+
+        case VelRepr.Mixed:
+            W_H_B = data.base_transform()
+            W_H_C = W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))  # noqa: F841
+            W_ṗ_B = data.base_velocity()[0:3]
+            W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)  # noqa: F841
 
         case _:
             raise ValueError(data.velocity_representation)
@@ -822,35 +1411,60 @@ def inverse_dynamics(
     # representation. In Mixed representation, this conversion is not a plain
     # transformation with just X, but it also involves a cross product in ℝ⁶.
     W_v̇_WB = to_inertial(
-        C_v̇_WB=base_acceleration,
+        C_v̇_WB=v̇_WB,
         W_H_C=W_H_C,
         C_v_WB=data.base_velocity(),
-        W_vl_WC=W_vl_WC,
+        W_v_WC=W_v_WC,
     )
 
+    # Create a references object that simplifies converting among representations.
     references = js.references.JaxSimModelReferences.build(
         model=model,
         data=data,
-        link_forces=external_forces,
+        link_forces=f_L,
         velocity_representation=data.velocity_representation,
     )
 
-    # Compute RNEA
+    # Extract the link and joint serializations.
+    link_names = model.link_names()
+    joint_names = model.joint_names()
+
+    # Extract the state in inertial-fixed representation.
+    with data.switch_velocity_representation(VelRepr.Inertial):
+        W_p_B = data.base_position()
+        W_v_WB = data.base_velocity()
+        W_Q_B = data.base_orientation(dcm=False)
+        s = data.joint_positions(model=model, joint_names=joint_names)
+        ṡ = data.joint_velocities(model=model, joint_names=joint_names)
+
+    # Extract the inputs in inertial-fixed representation.
     with references.switch_velocity_representation(VelRepr.Inertial):
-        W_f_B, τ = jaxsim.physics.algos.rnea.rnea(
-            model=model.physics_model,
-            xfb=data.state.physics_model.xfb(),
-            q=data.state.physics_model.joint_positions,
-            qd=data.state.physics_model.joint_velocities,
-            qdd=joint_accelerations,
-            a0fb=W_v̇_WB,
-            f_ext=references.link_forces(model=model, data=data),
-        )
+        W_f_L = references.link_forces(model=model, data=data, link_names=link_names)
+
+    # ========================
+    # Compute inverse dynamics
+    # ========================
+
+    W_f_B, τ = jaxsim.rbda.rnea(
+        model=model,
+        base_position=W_p_B,
+        base_quaternion=W_Q_B,
+        joint_positions=s,
+        base_linear_velocity=W_v_WB[0:3],
+        base_angular_velocity=W_v_WB[3:6],
+        joint_velocities=ṡ,
+        base_linear_acceleration=W_v̇_WB[0:3],
+        base_angular_acceleration=W_v̇_WB[3:6],
+        joint_accelerations=s̈,
+        link_forces=W_f_L,
+        standard_gravity=data.standard_gravity(),
+    )
 
-    # Adjust shape
-    τ = jnp.atleast_1d(τ.squeeze())
+    # =============
+    # Adjust output
+    # =============
 
-    # Express W_f_B in the active representation
+    # Express W_f_B in the active representation.
     f_B = js.data.JaxSimModelData.inertial_to_other_representation(
         array=W_f_B,
         other_representation=data.velocity_representation,
@@ -862,10 +1476,11 @@ def inverse_dynamics(
 
 
 @jax.jit
+@js.common.named_scope
 def free_floating_gravity_forces(
     model: JaxSimModel, data: js.data.JaxSimModelData
 ) -> jtp.Vector:
-    """
+    r"""
     Compute the free-floating gravity forces :math:`g(\mathbf{q})` of the model.
 
     Args:
@@ -876,12 +1491,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
     ):
@@ -905,16 +1520,17 @@ def free_floating_gravity_forces(
             # Set zero inputs:
             joint_accelerations=jnp.atleast_1d(jnp.zeros(model.dofs())),
             base_acceleration=jnp.zeros(6),
-            external_forces=jnp.zeros(shape=(model.number_of_links(), 6)),
+            link_forces=jnp.zeros(shape=(model.number_of_links(), 6)),
         )
     ).astype(float)
 
 
 @jax.jit
+@js.common.named_scope
 def free_floating_bias_forces(
     model: JaxSimModel, data: js.data.JaxSimModelData
 ) -> jtp.Vector:
-    """
+    r"""
     Compute the free-floating bias forces :math:`h(\mathbf{q}, \boldsymbol{\nu})`
     of the model.
 
@@ -926,12 +1542,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
     ):
@@ -948,18 +1564,20 @@ def free_floating_bias_forces(
             data.state.physics_model.joint_positions
         )
 
-        data_rnea.state.physics_model.base_linear_velocity = (
-            data.state.physics_model.base_linear_velocity
-        )
-
-        data_rnea.state.physics_model.base_angular_velocity = (
-            data.state.physics_model.base_angular_velocity
-        )
-
         data_rnea.state.physics_model.joint_velocities = (
             data.state.physics_model.joint_velocities
         )
 
+        # Make sure that base velocity is zero for fixed-base model.
+        if model.floating_base():
+            data_rnea.state.physics_model.base_linear_velocity = (
+                data.state.physics_model.base_linear_velocity
+            )
+
+            data_rnea.state.physics_model.base_angular_velocity = (
+                data.state.physics_model.base_angular_velocity
+            )
+
     return jnp.hstack(
         inverse_dynamics(
             model=model,
@@ -967,7 +1585,7 @@ def free_floating_bias_forces(
             # Set zero inputs:
             joint_accelerations=jnp.atleast_1d(jnp.zeros(model.dofs())),
             base_acceleration=jnp.zeros(6),
-            external_forces=jnp.zeros(shape=(model.number_of_links(), 6)),
+            link_forces=jnp.zeros(shape=(model.number_of_links(), 6)),
         )
     ).astype(float)
 
@@ -978,6 +1596,26 @@ def free_floating_bias_forces(
 
 
 @jax.jit
+@js.common.named_scope
+def locked_spatial_inertia(
+    model: JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    """
+    Compute the locked 6D inertia matrix of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The locked 6D inertia matrix of the model.
+    """
+
+    return total_momentum_jacobian(model=model, data=data)[:, 0:6]
+
+
+@jax.jit
+@js.common.named_scope
 def total_momentum(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vector:
     """
     Compute the total momentum of the model.
@@ -987,35 +1625,453 @@ def total_momentum(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vec
         data: The data of the considered model.
 
     Returns:
-        The total momentum of the model.
+        The total momentum of the model in the active velocity representation.
     """
 
-    # Compute the momentum in body-fixed velocity representation.
-    # Note: the first 6 rows of the mass matrix define the jacobian of the
-    #       floating-base momentum.
-    with data.switch_velocity_representation(velocity_representation=VelRepr.Body):
-        B_ν = data.generalized_velocity()
-        M_B = free_floating_mass_matrix(model=model, data=data)
+    ν = data.generalized_velocity()
+    Jh = total_momentum_jacobian(model=model, data=data)
+
+    return Jh @ ν
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def total_momentum_jacobian(
+    model: JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    """
+    Compute the jacobian of the total momentum.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        output_vel_repr: The output velocity representation of the jacobian.
+
+    Returns:
+        The jacobian of the total momentum of the model in the active representation.
+    """
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    if output_vel_repr is data.velocity_representation:
+        return free_floating_mass_matrix(model=model, data=data)[0:6]
+
+    with data.switch_velocity_representation(VelRepr.Body):
+        B_Jh_B = free_floating_mass_matrix(model=model, data=data)[0:6]
+
+    match data.velocity_representation:
+        case VelRepr.Body:
+            B_Jh = B_Jh_B
+
+        case VelRepr.Inertial:
+            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 = 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 _:
+            raise ValueError(data.velocity_representation)
+
+    match output_vel_repr:
+        case VelRepr.Body:
+            return B_Jh
+
+        case VelRepr.Inertial:
+            W_H_B = data.base_transform()
+            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 = 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
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+
+@jax.jit
+@js.common.named_scope
+def average_velocity(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vector:
+    """
+    Compute the average velocity of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The average velocity of the model computed in the base frame and expressed
+        in the active representation.
+    """
+
+    ν = data.generalized_velocity()
+    J = average_velocity_jacobian(model=model, data=data)
+
+    return J @ ν
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def average_velocity_jacobian(
+    model: JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    """
+    Compute the Jacobian of the average velocity of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        output_vel_repr: The output velocity representation of the jacobian.
+
+    Returns:
+        The Jacobian of the average centroidal velocity of the model in the desired
+        representation.
+    """
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Depending on the velocity representation, the frame G is either G[W] or G[B].
+    G_J = js.com.average_centroidal_velocity_jacobian(model=model, data=data)
+
+    match output_vel_repr:
+
+        case VelRepr.Inertial:
+
+            GW_J = G_J
+            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 = Adjoint.from_transform(transform=W_H_GW)
+
+            return W_X_GW @ GW_J
+
+        case VelRepr.Body:
+
+            GB_J = G_J
+            W_p_B = data.base_position()
+            W_p_CoM = js.com.com_position(model=model, data=data)
+            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 = Adjoint.from_transform(transform=B_H_GB)
+
+            return B_X_GB @ GB_J
+
+        case VelRepr.Mixed:
 
-    # Compute the total momentum expressed in the base frame
-    B_h = M_B[0:6, :] @ B_ν
+            GW_J = G_J
+            W_p_B = data.base_position()
+            W_p_CoM = js.com.com_position(model=model, data=data)
 
-    # Compute the 6D transformation matrix
+            BW_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM - W_p_B)
+            BW_X_GW = Adjoint.from_transform(transform=BW_H_GW)
+
+            return BW_X_GW @ GW_J
+
+
+# ========================
+# Other dynamic quantities
+# ========================
+
+
+@jax.jit
+@js.common.named_scope
+def link_bias_accelerations(
+    model: JaxSimModel,
+    data: js.data.JaxSimModelData,
+) -> jtp.Vector:
+    r"""
+    Compute the bias accelerations of the links of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The bias accelerations of the links of the model.
+
+    Note:
+        This function computes the component of the total 6D acceleration not due to
+        the joint or base acceleration.
+        It is often called :math:`\dot{J} \boldsymbol{\nu}`.
+    """
+
+    # ================================================
+    # Compute the body-fixed zero base 6D acceleration
+    # ================================================
+
+    # Compute the base transform.
     W_H_B = data.base_transform()
-    B_X_W: jtp.Array = jaxlie.SE3.from_matrix(W_H_B).inverse().adjoint()
 
-    # Convert to inertial-fixed representation
-    # (its coordinates transform like 6D forces)
-    W_h = B_X_W.T @ B_h
+    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
+    ) -> jtp.Vector:
+        """
+        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.
+        """
+
+        W_X_C = Adjoint.from_transform(transform=W_H_C)
+        C_X_W = Adjoint.from_transform(transform=W_H_C, inverse=True)
 
-    # Convert to the active representation of the model
-    return js.data.JaxSimModelData.inertial_to_other_representation(
-        array=W_h,
-        other_representation=data.velocity_representation,
-        transform=W_H_B,
-        is_force=True,
+        # In Mixed representation, we need to include a cross product in ℝ⁶.
+        # In Inertial and Body representations, the cross product is always zero.
+        return W_X_C @ (C_v̇_WB + jaxsim.math.Cross.vx(C_X_W @ W_v_WC) @ C_v_WB)
+
+    # Here we initialize a zero 6D acceleration in the active representation, and
+    # convert it to inertial-fixed. This is a useful intermediate representation
+    # because the apparent acceleration W_v̇_WB is equal to the intrinsic acceleration
+    # W_a_WB, and intrinsic accelerations can be expressed in different frames through
+    # a simple C_X_W 6D transform.
+    match data.velocity_representation:
+        case VelRepr.Inertial:
+            W_H_C = W_H_W = jnp.eye(4)  # noqa: F841
+            W_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                C_v_WB = W_v_WB = data.base_velocity()
+
+        case VelRepr.Body:
+            W_H_C = W_H_B
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                W_v_WC = W_v_WB = data.base_velocity()  # noqa: F841
+            with data.switch_velocity_representation(VelRepr.Body):
+                C_v_WB = B_v_WB = data.base_velocity()
+
+        case VelRepr.Mixed:
+            W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
+            W_H_C = W_H_BW
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                W_ṗ_B = data.base_velocity()[0:3]
+                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  # noqa: F841
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                C_v_WB = BW_v_WB = data.base_velocity()  # noqa: F841
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Convert a zero 6D acceleration from the active representation to inertial-fixed.
+    W_v̇_WB = other_representation_to_inertial(
+        C_v̇_WB=jnp.zeros(6), C_v_WB=C_v_WB, W_H_C=W_H_C, W_v_WC=W_v_WC
+    )
+
+    # ===================================
+    # Initialize buffers and prepare data
+    # ===================================
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute 6D transforms of the base velocity.
+    B_X_W = jaxsim.math.Adjoint.from_transform(transform=W_H_B, inverse=True)
+
+    # 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(
+        joint_positions=data.joint_positions(), base_transform=W_H_B
+    )
+
+    # Allocate the buffer to store the body-fixed link velocities.
+    L_v_WL = jnp.zeros(shape=(model.number_of_links(), 6))
+
+    # Store the base velocity.
+    with data.switch_velocity_representation(VelRepr.Body):
+        B_v_WB = data.base_velocity()
+        L_v_WL = L_v_WL.at[0].set(B_v_WB)
+
+    # Get the joint velocities.
+    ṡ = data.joint_velocities(model=model, joint_names=model.joint_names())
+
+    # Allocate the buffer to store the body-fixed link accelerations,
+    # and initialize the base acceleration.
+    L_v̇_WL = jnp.zeros(shape=(model.number_of_links(), 6))
+    L_v̇_WL = L_v̇_WL.at[0].set(B_X_W @ W_v̇_WB)
+
+    # ======================================
+    # Propagate accelerations and velocities
+    # ======================================
+
+    # The computation of the bias forces is similar to the forward pass of RNEA,
+    # this time with zero base and joint accelerations. Furthermore, here we do
+    # not remove gravity during the propagation.
+
+    # Initialize the loop.
+    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]:
+        # Initialize index and unpack the carry.
+        ii = i - 1
+        v, a = carry
+
+        # Get the motion subspace of the joint.
+        Si = S[i].squeeze()
+
+        # Project the joint velocity into its motion subspace.
+        vJ = Si * ṡ[ii]
+
+        # Propagate the link body-fixed velocity.
+        v_i = i_X_λi[i] @ v[λ[i]] + vJ
+        v = v.at[i].set(v_i)
+
+        # Propagate the link body-fixed acceleration considering zero joint acceleration.
+        s̈ = 0.0
+        a_i = i_X_λi[i] @ a[λ[i]] + Si * s̈ + jaxsim.math.Cross.vx(v[i]) @ vJ
+        a = a.at[i].set(a_i)
+
+        return (v, a), None
+
+    # Compute the body-fixed velocity and body-fixed apparent acceleration of the links.
+    (L_v_WL, L_v̇_WL), _ = (
+        jax.lax.scan(
+            f=propagate_accelerations,
+            init=carry0,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(L_v_WL, L_v̇_WL), None]
+    )
+
+    # ===================================================================
+    # Convert the body-fixed 6D acceleration to the active representation
+    # ===================================================================
+
+    def body_to_other_representation(
+        L_v̇_WL: jtp.Vector, L_v_WL: jtp.Vector, C_H_L: jtp.Matrix, L_v_CL: jtp.Vector
+    ) -> jtp.Vector:
+        """
+        Convert the body-fixed apparent acceleration L_v̇_WL to
+        another representation C_v̇_WL expressed in a generic frame C.
+        """
+
+        # In Mixed representation, we need to include a cross product in ℝ⁶.
+        # In Inertial and Body representations, the cross product is always zero.
+        C_X_L = jaxsim.math.Adjoint.from_transform(transform=C_H_L)
+        return C_X_L @ (L_v̇_WL + jaxsim.math.Cross.vx(L_v_CL) @ L_v_WL)
+
+    match data.velocity_representation:
+        case VelRepr.Body:
+            C_H_L = L_H_L = jnp.stack(  # noqa: F841
+                [jnp.eye(4)] * model.number_of_links()
+            )
+            L_v_CL = L_v_LL = jnp.zeros(  # noqa: F841
+                shape=(model.number_of_links(), 6)
+            )
+
+        case VelRepr.Inertial:
+            C_H_L = W_H_L = js.model.forward_kinematics(model=model, data=data)
+            L_v_CL = L_v_WL
+
+        case VelRepr.Mixed:
+            W_H_L = js.model.forward_kinematics(model=model, data=data)
+            LW_H_L = jax.vmap(lambda W_H_L: W_H_L.at[0:3, 3].set(jnp.zeros(3)))(W_H_L)
+            C_H_L = LW_H_L
+            L_v_CL = L_v_LW_L = jax.vmap(  # noqa: F841
+                lambda v: v.at[0:3].set(jnp.zeros(3))
+            )(L_v_WL)
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Convert from body-fixed to the active representation.
+    O_v̇_WL = jax.vmap(body_to_other_representation)(
+        L_v̇_WL=L_v̇_WL, L_v_WL=L_v_WL, C_H_L=C_H_L, L_v_CL=L_v_CL
+    )
+
+    return O_v̇_WL
+
+
+@jax.jit
+@js.common.named_scope
+def link_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
+    **kwargs,
+) -> jtp.Matrix:
+    """
+    Compute the 6D contact forces of all links of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_forces:
+            The 6D external forces to apply to the links expressed in the same
+            representation of data.
+        joint_force_references:
+            The joint force references to apply to the joints.
+        kwargs: Additional keyword arguments to pass to the active contact model..
+
+    Returns:
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
+        expressed in the frame corresponding to the active representation.
+    """
+
+    # Note: the following code should be kept in sync with the function
+    # `jaxsim.api.ode.system_velocity_dynamics`. We cannot merge them since
+    # there we need to get also aux_data.
+
+    # 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
+        else jnp.zeros((model.number_of_links(), 6))
     ).astype(float)
 
+    # Build joint force references if not provided.
+    joint_force_references = (
+        jnp.atleast_1d(joint_force_references)
+        if joint_force_references is not None
+        else jnp.zeros(model.dofs())
+    )
+
+    # We expect that the 6D forces included in the `link_forces` argument are expressed
+    # in the frame corresponding to the velocity representation of `data`.
+    input_references = js.references.JaxSimModelReferences.build(
+        model=model,
+        data=data,
+        velocity_representation=data.velocity_representation,
+        link_forces=O_f_L,
+        joint_force_references=joint_force_references,
+    )
+
+    # Compute the 6D forces applied to the links equivalent to the forces applied
+    # to the frames associated to the collidable points.
+    f_L, _ = model.contact_model.compute_link_contact_forces(
+        model=model,
+        data=data,
+        link_forces=input_references.link_forces(model=model, data=data),
+        joint_force_references=input_references.joint_force_references(),
+        **kwargs,
+    )
+
+    return f_L
+
 
 # ======
 # Energy
@@ -1023,6 +2079,7 @@ def total_momentum(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vec
 
 
 @jax.jit
+@js.common.named_scope
 def mechanical_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
     Compute the mechanical energy of the model.
@@ -1042,6 +2099,7 @@ def mechanical_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.
 
 
 @jax.jit
+@js.common.named_scope
 def kinetic_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
     Compute the kinetic energy of the model.
@@ -1063,6 +2121,7 @@ def kinetic_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Flo
 
 
 @jax.jit
+@js.common.named_scope
 def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
     Compute the potential energy of the model.
@@ -1077,7 +2136,7 @@ def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.F
 
     m = total_mass(model=model)
     gravity = data.gravity.squeeze()
-    W_p̃_CoM = jnp.hstack([com_position(model=model, data=data), 1])
+    W_p̃_CoM = jnp.hstack([js.com.com_position(model=model, data=data), 1])
 
     U = -jnp.hstack([gravity, 0]) @ (m * W_p̃_CoM)
     return U.squeeze().astype(float)
@@ -1089,15 +2148,18 @@ def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.F
 
 
 @jax.jit
+@js.common.named_scope
 def step(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    dt: jtp.FloatLike,
-    integrator: jaxsim.integrators.Integrator,
-    integrator_state: dict[str, Any] | None = None,
-    joint_forces: jtp.VectorLike | None = None,
-    external_forces: jtp.MatrixLike | None = None,
+    t0: jtp.FloatLike = 0.0,
+    dt: jtp.FloatLike | None = None,
+    integrator: jaxsim.integrators.Integrator | None = None,
+    integrator_metadata: dict[str, Any] | None = None,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
+    **kwargs,
 ) -> tuple[js.data.JaxSimModelData, dict[str, Any]]:
     """
     Perform a simulation step.
@@ -1105,40 +2167,188 @@ def step(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        dt: The time step to consider.
         integrator: The integrator to use.
-        integrator_state: The state of the integrator.
-        joint_forces: The joint forces to consider.
-        external_forces:
-            The external forces to consider.
-            The frame in which they are expressed must be `data.velocity_representation`.
+        integrator_metadata: The metadata of the integrator, if needed.
+        t0: The initial time to consider. Only relevant for time-dependent dynamics.
+        dt: The time step to consider. If not specified, it is read from the model.
+        link_forces:
+            The 6D forces to apply to the links expressed in the frame corresponding to
+            the velocity representation of `data`.
+        joint_force_references: The joint force references to consider.
+        kwargs: Additional kwargs to pass to the integrator.
 
     Returns:
-        A tuple containing the new data of the model
-        and the new state of the integrator.
+        A tuple containing the new data of the model and a dictionary of auxiliary
+        data computed during the step. If the integrator has metadata, the dictionary
+        will contain the new metadata stored in the `integrator_metadata` key.
+
+    Note:
+        In order to reduce the occurrences of frame conversions performed internally,
+        it is recommended to use inertial-fixed velocity representation. This can be
+        particularly useful for automatically differentiated logic.
     """
 
-    integrator_state = integrator_state if integrator_state is not None else dict()
+    # Extract the integrator kwargs.
+    # The following logic allows using integrators having kwargs colliding with the
+    # kwargs of this step function.
+    kwargs = kwargs if kwargs is not None else {}
+    integrator_kwargs = kwargs.pop("integrator_kwargs", {})
+    integrator_kwargs = kwargs | integrator_kwargs
+
+    # Extract the integrator and the optional metadata.
+    integrator_metadata_t0 = integrator_metadata
+    integrator = integrator if integrator is not None else model.integrator
+
+    # Initialize the time-related variables.
+    state_t0 = data.state
+    t0 = jnp.array(t0, dtype=float)
+    dt = jnp.array(dt if dt is not None else model.time_step).astype(float)
+
+    # The visco-elastic contacts operate at best with their own integrator.
+    # They can be used with Euler-like integrators, paying the price of ignoring
+    # some of the benefits of continuous-time integration on the system position.
+    # Furthermore, the requirement to know the Δt used by the integrator is not
+    # compatible with high-order integrators, that use advanced RK stages to evaluate
+    # the dynamics at intermediate times.
+    module = jaxsim.rbda.contacts.visco_elastic.step.__module__
+    name = jaxsim.rbda.contacts.visco_elastic.step.__name__
+    msg = "You need to use the custom '{}.{}' function with this contact model."
+    jaxsim.exceptions.raise_runtime_error_if(
+        condition=(
+            isinstance(model.contact_model, jaxsim.rbda.contacts.ViscoElasticContacts)
+            & (
+                ~jnp.allclose(dt, model.time_step)
+                | ~int(
+                    isinstance(integrator, jaxsim.integrators.fixed_step.ForwardEuler)
+                )
+            )
+        ),
+        msg=msg.format(module, name),
+    )
 
-    # Extract the initial resources.
-    t0_ns = data.time_ns
-    state_x0 = data.state
-    integrator_state_x0 = integrator_state
+    # =================
+    # Phase 1: pre-step
+    # =================
+
+    # TODO: some contact models here may want to perform a dynamic filtering of
+    # the enabled collidable points.
+
+    # Build the references object.
+    # We assume that the link forces are expressed in the frame corresponding to the
+    # velocity representation of the data.
+    references = js.references.JaxSimModelReferences.build(
+        model=model,
+        data=data,
+        velocity_representation=data.velocity_representation,
+        link_forces=link_forces,
+        joint_force_references=joint_force_references,
+    )
+
+    # =============
+    # Phase 2: step
+    # =============
+
+    # Prepare the references to pass.
+    with references.switch_velocity_representation(data.velocity_representation):
+
+        f_L = references.link_forces(model=model, data=data)
+        τ_references = references.joint_force_references(model=model)
 
     # Step the dynamics forward.
-    state_xf, integrator_state_xf = integrator.step(
-        x0=state_x0,
-        t0=jnp.array(t0_ns * 1e9).astype(float),
+    state_tf, integrator_metadata_tf = integrator.step(
+        x0=state_t0,
+        t0=t0,
         dt=dt,
-        params=integrator_state_x0,
-        **dict(joint_forces=joint_forces, external_forces=external_forces),
+        metadata=integrator_metadata_t0,
+        # Always inject the current (model, data) pair into the system dynamics
+        # considered by the integrator, and include the input variables represented
+        # by the pair (f_L, τ_references).
+        # Note that the wrapper of the system dynamics will override (state_x0, t0)
+        # inside the passed data even if it is not strictly needed. This logic is
+        # necessary to reuse the jit-compiled step function of compatible pytrees
+        # of model and data produced e.g. by parameterized applications.
+        **(
+            dict(
+                model=model,
+                data=data,
+                link_forces=f_L,
+                joint_force_references=τ_references,
+            )
+            | integrator_kwargs
+        ),
     )
 
-    return (
-        # Store the new state of the model and the new time.
-        data.replace(
-            state=state_xf,
-            time_ns=t0_ns + jnp.array(dt * 1e9).astype(jnp.uint64),
-        ),
-        integrator_state_xf,
+    # Store the new state of the model.
+    data_tf = data.replace(state=state_tf)
+
+    # ==================
+    # Phase 3: post-step
+    # ==================
+
+    # Post process the simulation state, if needed.
+    match model.contact_model:
+
+        # Rigid contact models use an impact model that produces discontinuous model velocities.
+        # Hence, here we need to reset the velocity after each impact to guarantee that
+        # the linear velocity of the active collidable points is zero.
+        case jaxsim.rbda.contacts.RigidContacts():
+
+            # Raise runtime error for not supported case in which Rigid contacts and
+            # Baumgarte stabilization are enabled and used with ForwardEuler integrator.
+            jaxsim.exceptions.raise_runtime_error_if(
+                condition=isinstance(
+                    integrator,
+                    jaxsim.integrators.fixed_step.ForwardEuler
+                    | jaxsim.integrators.fixed_step.ForwardEulerSO3,
+                )
+                & ((data_tf.contacts_params.K > 0) | (data_tf.contacts_params.D > 0)),
+                msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
+            )
+
+            # Extract the indices corresponding to the enabled collidable points.
+            indices_of_enabled_collidable_points = (
+                model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+            )
+
+            W_p_C = js.contact.collidable_point_positions(model, data_tf)[
+                indices_of_enabled_collidable_points
+            ]
+
+            # Compute the penetration depth of the collidable points.
+            δ, *_ = jax.vmap(
+                jaxsim.rbda.contacts.common.compute_penetration_data,
+                in_axes=(0, 0, None),
+            )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
+
+            with data_tf.switch_velocity_representation(VelRepr.Mixed):
+                J_WC = js.contact.jacobian(model, data_tf)[
+                    indices_of_enabled_collidable_points
+                ]
+                M = js.model.free_floating_mass_matrix(model, data_tf)
+                BW_ν_pre_impact = data_tf.generalized_velocity()
+
+                # Compute the impact velocity.
+                # It may be discontinuous in case new contacts are made.
+                BW_ν_post_impact = (
+                    jaxsim.rbda.contacts.RigidContacts.compute_impact_velocity(
+                        generalized_velocity=BW_ν_pre_impact,
+                        inactive_collidable_points=(δ <= 0),
+                        M=M,
+                        J_WC=J_WC,
+                    )
+                )
+
+                # Reset the generalized velocity.
+                data_tf = data_tf.reset_base_velocity(BW_ν_post_impact[0:6])
+                data_tf = data_tf.reset_joint_velocities(BW_ν_post_impact[6:])
+
+    # Restore the input velocity representation.
+    data_tf = data_tf.replace(
+        velocity_representation=data.velocity_representation, validate=False
+    )
+
+    return data_tf, {} | (
+        dict(integrator_metadata=integrator_metadata_tf)
+        if integrator_metadata is not None
+        else {}
     )