jaxsim 0.4.1.dev26__py3-none-any.whl → 0.4.2__py3-none-any.whl

jaxsim/api/model.py

@@ -301,10 +301,10 @@ class JaxSimModel(JaxsimDataclass):
 
     def frame_names(self) -> tuple[str, ...]:
         """
-        Return the names of the links in the model.
+        Return the names of the frames in the model.
 
         Returns:
-            The names of the links in the model.
+            The names of the frames in the model.
         """
 
         return self.kin_dyn_parameters.frame_parameters.name
@@ -495,8 +495,9 @@ def generalized_free_floating_jacobian(
             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 = B_J_full_WX_B @ jax.scipy.linalg.block_diag(
-                B_X_W, jnp.eye(model.dofs())
+            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:
@@ -509,7 +510,7 @@ def generalized_free_floating_jacobian(
             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 = (
+            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()))
             )
@@ -542,11 +543,13 @@ def generalized_free_floating_jacobian(
             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(lambda B_J_WL_I: W_X_B @ B_J_WL_I)(B_J_WL_I)
+            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:
 
-            O_J_WL_I = L_J_WL_I = jax.vmap(
+            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
                 )
@@ -565,7 +568,7 @@ def generalized_free_floating_jacobian(
                 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(
+            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)
@@ -576,6 +579,41 @@ def generalized_free_floating_jacobian(
     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
+    )
+
+    O_J̇_WL_I = jax.vmap(
+        lambda model, data, link_idxs, output_vel_repr: js.link.jacobian_derivative(
+            model, data, link_index=link_idxs, output_vel_repr=output_vel_repr
+        ),
+        in_axes=(None, None, 0, None),
+    )(model, data, jnp.arange(model.number_of_links()), output_vel_repr)
+
+    return O_J̇_WL_I
+
+
 @functools.partial(jax.jit, static_argnames=["prefer_aba"])
 def forward_dynamics(
     model: JaxSimModel,
@@ -721,8 +759,8 @@ def forward_dynamics_aba(
     match data.velocity_representation:
         case VelRepr.Inertial:
             # In this case C=W
-            W_H_C = W_H_W = jnp.eye(4)
-            W_v_WC = W_v_WW = jnp.zeros(6)
+            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
@@ -732,9 +770,9 @@ def forward_dynamics_aba(
         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_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)
+            W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)  # noqa: F841
 
         case _:
             raise ValueError(data.velocity_representation)
@@ -1089,8 +1127,8 @@ def inverse_dynamics(
 
     match data.velocity_representation:
         case VelRepr.Inertial:
-            W_H_C = W_H_W = jnp.eye(4)
-            W_v_WC = W_v_WW = jnp.zeros(6)
+            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()
@@ -1099,9 +1137,9 @@ def inverse_dynamics(
 
         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_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)
+            W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)  # noqa: F841
 
         case _:
             raise ValueError(data.velocity_representation)
@@ -1536,15 +1574,15 @@ def link_bias_accelerations(
     # a simple C_X_W 6D transform.
     match data.velocity_representation:
         case VelRepr.Inertial:
-            W_H_C = W_H_W = jnp.eye(4)
-            W_v_WC = W_v_WW = jnp.zeros(6)
+            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()
+                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()
 
@@ -1555,9 +1593,9 @@ def link_bias_accelerations(
                 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
+                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()
+                C_v_WB = BW_v_WB = data.base_velocity()  # noqa: F841
 
         case _:
             raise ValueError(data.velocity_representation)
@@ -1665,8 +1703,12 @@ def link_bias_accelerations(
 
     match data.velocity_representation:
         case VelRepr.Body:
-            C_H_L = L_H_L = jnp.stack([jnp.eye(4)] * model.number_of_links())
-            L_v_CL = L_v_LL = jnp.zeros(shape=(model.number_of_links(), 6))
+            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)
@@ -1676,7 +1718,9 @@ def link_bias_accelerations(
             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(lambda v: v.at[0:3].set(jnp.zeros(3)))(L_v_WL)
+            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)

jaxsim/api/references.py

@@ -8,6 +8,7 @@ import jax_dataclasses
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim import exceptions
 from jaxsim.utils.tracing import not_tracing
 
 from .common import VelRepr
@@ -30,6 +31,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
     @staticmethod
     def zero(
         model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData | None = None,
         velocity_representation: VelRepr = VelRepr.Inertial,
     ) -> JaxSimModelReferences:
         """
@@ -37,6 +39,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         Args:
             model: The model for which to create the zero references.
+            data:
+                The data of the model, only needed if the velocity representation is
+                not inertial-fixed.
             velocity_representation: The velocity representation to use.
 
         Returns:
@@ -44,7 +49,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         """
 
         return JaxSimModelReferences.build(
-            model=model, velocity_representation=velocity_representation
+            model=model, data=data, velocity_representation=velocity_representation
         )
 
     @staticmethod
@@ -441,3 +446,104 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         return replace(
             forces=self.input.physics_model.f_ext.at[link_idxs, :].set(W_f0_L + W_f_L)
         )
+
+    def apply_frame_forces(
+        self,
+        forces: jtp.MatrixLike,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        frame_names: tuple[str, ...] | str | None = None,
+        additive: bool = False,
+    ) -> Self:
+        """
+        Apply the frame forces.
+
+        Args:
+            forces: The frame 6D forces in the active representation.
+            model:
+                The model to consider, only needed if a frame serialization different
+                from the implicit one is used.
+            data:
+                The data of the considered model, only needed if the velocity
+                representation is not inertial-fixed.
+            frame_names: The names of the frames corresponding to the forces.
+            additive:
+                Whether to add the forces to the existing ones instead of replacing them.
+
+        Returns:
+            A new `JaxSimModelReferences` object with the given frame forces.
+
+        Note:
+            The frame forces must be expressed in the active representation.
+            Then, we always convert and store forces in inertial-fixed representation.
+        """
+
+        f_F = jnp.atleast_2d(forces).astype(float)
+
+        # If we have the model, we can extract the frame names if not provided.
+        frame_names = frame_names if frame_names is not None else model.frame_names()
+
+        # Make sure that the frame names are a tuple if they are provided by the user.
+        frame_names = (frame_names,) if isinstance(frame_names, str) else frame_names
+
+        if len(frame_names) != f_F.shape[0]:
+            msg = "The number of frame names ({}) must match the number of forces ({})"
+            raise ValueError(msg.format(len(frame_names), f_F.shape[0]))
+
+        # Extract the frame indices.
+        frame_idxs = js.frame.names_to_idxs(frame_names=frame_names, model=model)
+        parent_link_idxs = jax.vmap(js.frame.idx_of_parent_link, in_axes=(None,))(
+            model, frame_index=frame_idxs
+        )
+
+        exceptions.raise_value_error_if(
+            condition=jnp.logical_not(data.valid(model=model)),
+            msg="The provided data is not valid for the model",
+        )
+        W_H_Fi = jax.vmap(
+            lambda frame_idx: js.frame.transform(
+                model=model, data=data, frame_index=frame_idx
+            )
+        )(frame_idxs)
+
+        # Helper function to convert a single 6D force to the inertial representation
+        # considering as body the frame (i.e. L_f_F and LW_f_F).
+        def to_inertial(f_F: jtp.MatrixLike, W_H_F: jtp.MatrixLike) -> jtp.Matrix:
+            return JaxSimModelReferences.other_representation_to_inertial(
+                array=f_F,
+                other_representation=self.velocity_representation,
+                transform=W_H_F,
+                is_force=True,
+            )
+
+        match self.velocity_representation:
+            case VelRepr.Inertial:
+                W_f_F = f_F
+
+            case VelRepr.Body | VelRepr.Mixed:
+                W_f_F = jax.vmap(to_inertial)(f_F, W_H_Fi)
+
+            case _:
+                raise ValueError("Invalid velocity representation.")
+
+        # Sum the forces on the parent links.
+        mask = parent_link_idxs[:, jnp.newaxis] == jnp.arange(model.number_of_links())
+        W_f_L = mask.T @ W_f_F
+
+        with self.switch_velocity_representation(
+            velocity_representation=VelRepr.Inertial
+        ):
+            references = self.apply_link_forces(
+                model=model,
+                data=data,
+                link_names=js.link.idxs_to_names(
+                    model=model, link_indices=parent_link_idxs
+                ),
+                forces=W_f_L,
+                additive=additive,
+            )
+
+        with references.switch_velocity_representation(
+            velocity_representation=self.velocity_representation
+        ):
+            return references

jaxsim/integrators/common.py

@@ -261,8 +261,10 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         # Check if the Butcher tableau supports FSAL (first-same-as-last).
         # If it does, store the index of the intermediate derivative to be used as the
         # first derivative of the next iteration.
-        has_fsal, index_of_fsal = ExplicitRungeKutta.butcher_tableau_supports_fsal(
-            A=cls.A, b=cls.b, c=cls.c, index_of_solution=cls.row_index_of_solution
+        has_fsal, index_of_fsal = (  # noqa: F841
+            ExplicitRungeKutta.butcher_tableau_supports_fsal(
+                A=cls.A, b=cls.b, c=cls.c, index_of_solution=cls.row_index_of_solution
+            )
         )
 
         # Build the integrator object.

jaxsim/math/adjoint.py

@@ -3,7 +3,6 @@ import jaxlie
 
 import jaxsim.typing as jtp
 
-from .quaternion import Quaternion
 from .skew import Skew
 
 
@@ -31,7 +30,7 @@ class Adjoint:
         assert quaternion.size == 4
         assert translation.size == 3
 
-        Q_sixd = jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(quaternion))
+        Q_sixd = jaxlie.SO3(wxyz=quaternion)
         Q_sixd = Q_sixd if not normalize_quaternion else Q_sixd.normalize()
 
         return Adjoint.from_rotation_and_translation(
@@ -84,14 +83,14 @@ class Adjoint:
         A_o_B = translation.squeeze()
 
         if not inverse:
-            X = A_X_B = jnp.vstack(
+            X = A_X_B = jnp.vstack(  # noqa: F841
                 [
                     jnp.block([A_R_B, Skew.wedge(A_o_B) @ A_R_B]),
                     jnp.block([jnp.zeros(shape=(3, 3)), A_R_B]),
                 ]
             )
         else:
-            X = B_X_A = jnp.vstack(
+            X = B_X_A = jnp.vstack(  # noqa: F841
                 [
                     jnp.block([A_R_B.T, -A_R_B.T @ Skew.wedge(A_o_B)]),
                     jnp.block([jnp.zeros(shape=(3, 3)), A_R_B.T]),

jaxsim/math/joint_model.py

@@ -254,8 +254,8 @@ def supported_joint_motion(
         # This is a metadata required by only some joint types.
         axis = jnp.array(joint_axis).astype(float).squeeze()
 
-        pre_H_suc = jaxlie.SE3.from_rotation(
-            rotation=jaxlie.SO3.from_matrix(Rotation.from_axis_angle(vector=s * axis))
+        pre_H_suc = jaxlie.SE3.from_matrix(
+            matrix=jnp.eye(4).at[:3, :3].set(Rotation.from_axis_angle(vector=s * axis))
         )
 
         S = jnp.vstack(jnp.hstack([jnp.zeros(3), axis]))

jaxsim/math/quaternion.py

@@ -43,9 +43,7 @@ class Quaternion:
         Returns:
             jtp.Matrix: Direction cosine matrix (DCM).
         """
-        return jaxlie.SO3.from_quaternion_xyzw(
-            xyzw=Quaternion.to_xyzw(quaternion)
-        ).as_matrix()
+        return jaxlie.SO3(wxyz=quaternion).as_matrix()
 
     @staticmethod
     def from_dcm(dcm: jtp.Matrix) -> jtp.Vector:
@@ -158,7 +156,7 @@ class Quaternion:
         A_Q_B = jnp.array(quaternion).squeeze().astype(float)
 
         # Build the initial SO(3) quaternion.
-        W_Q_B_t0 = jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(wxyz=A_Q_B))
+        W_Q_B_t0 = jaxlie.SO3(wxyz=A_Q_B)
 
         # Integrate the quaternion on the manifold.
         W_Q_B_tf = jax.lax.select(

jaxsim/math/transform.py

@@ -3,8 +3,6 @@ import jaxlie
 
 import jaxsim.typing as jtp
 
-from .quaternion import Quaternion
-
 
 class Transform:
 
@@ -35,7 +33,7 @@ class Transform:
         assert W_p_B.size == 3
         assert W_Q_B.size == 4
 
-        A_R_B = jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(W_Q_B))
+        A_R_B = jaxlie.SO3(wxyz=W_Q_B)
         A_R_B = A_R_B if not normalize_quaternion else A_R_B.normalize()
 
         A_H_B = jaxlie.SE3.from_rotation_and_translation(

jaxsim/mujoco/loaders.py

@@ -532,7 +532,12 @@ class UrdfToMjcf:
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
-        cameras: list[dict[str, str]] | dict[str, str] | None = None,
+        cameras: (
+            MujocoCamera
+            | Sequence[MujocoCamera]
+            | dict[str, str]
+            | Sequence[dict[str, str]]
+        ) = (),
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a URDF file to a Mujoco MJCF string.
@@ -574,7 +579,12 @@ class SdfToMjcf:
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
-        cameras: list[dict[str, str]] | dict[str, str] | None = None,
+        cameras: (
+            MujocoCamera
+            | Sequence[MujocoCamera]
+            | dict[str, str]
+            | Sequence[dict[str, str]]
+        ) = (),
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a SDF file to a Mujoco MJCF string.

jaxsim/mujoco/model.py

@@ -238,9 +238,9 @@ class MujocoModelHelper:
             raise ValueError("The orientation is not a valid element of SO(3)")
 
         W_Q_B = (
-            Rotation.from_matrix(orientation).as_quat(canonical=True)[
-                np.array([3, 0, 1, 2])
-            ]
+            Rotation.from_matrix(orientation).as_quat(
+                canonical=True, scalar_first=False
+            )
             if dcm
             else orientation
         )
@@ -394,8 +394,8 @@ class MujocoModelHelper:
         if dcm:
             return R
 
-        q_xyzw = Rotation.from_matrix(R).as_quat(canonical=True)
-        return q_xyzw[[3, 0, 1, 2]]
+        q_xyzw = Rotation.from_matrix(R).as_quat(canonical=True, scalar_first=False)
+        return q_xyzw
 
     # ===============
     # Private methods

jaxsim/rbda/aba.py

@@ -4,7 +4,7 @@ import jaxlie
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint, Cross, Quaternion, StandardGravity
+from jaxsim.math import Adjoint, Cross, StandardGravity
 
 from . import utils
 
@@ -77,7 +77,7 @@ def aba(
 
     # Compute the base transform.
     W_H_B = jaxlie.SE3.from_rotation_and_translation(
-        rotation=jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(wxyz=W_Q_B)),
+        rotation=jaxlie.SO3(wxyz=W_Q_B),
         translation=W_p_B,
     )
 

jaxsim/rbda/collidable_points.py

@@ -4,7 +4,7 @@ import jaxlie
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint, Quaternion, Skew
+from jaxsim.math import Adjoint, Skew
 
 from . import utils
 
@@ -57,7 +57,7 @@ def collidable_points_pos_vel(
 
     # Compute the base transform.
     W_H_B = jaxlie.SE3.from_rotation_and_translation(
-        rotation=jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(wxyz=W_Q_B)),
+        rotation=jaxlie.SO3(wxyz=W_Q_B),
         translation=W_p_B,
     )
 

jaxsim/rbda/contacts/soft.py

@@ -192,7 +192,7 @@ class SoftContacts(ContactModel):
 
         # Unpack the position of the collidable point.
         px, py, pz = W_p_C = position.squeeze()
-        vx, vy, vz = W_ṗ_C = velocity.squeeze()
+        W_ṗ_C = velocity.squeeze()
 
         # Compute the terrain normal and the contact depth.
         n̂ = self.terrain.normal(x=px, y=py).squeeze()

jaxsim/rbda/crba.py

@@ -59,10 +59,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
         return (i_X_0,), None
 
-    (i_X_0,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=forward_pass_carry,
-        xs=jnp.arange(start=1, stop=model.number_of_links()),
+    (i_X_0,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=forward_pass_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(i_X_0,), None]
     )
 
     # ===================
@@ -128,10 +132,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
                 operand=carry,
             )
 
-        (j, Fi, M), _ = jax.lax.scan(
-            f=inner_fn,
-            init=carry_inner_fn,
-            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        (j, Fi, M), _ = (
+            jax.lax.scan(
+                f=inner_fn,
+                init=carry_inner_fn,
+                xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+            )
+            if model.number_of_links() > 1
+            else [(j, Fi, M), None]
         )
 
         Fi = i_X_0[j].T @ Fi
@@ -143,10 +151,14 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
     # This scan performs the backward pass to compute Mbj, Mjb and Mjj, that
     # also includes a fake while loop implemented with a scan and two cond.
-    (Mc, M), _ = jax.lax.scan(
-        f=backward_pass,
-        init=backward_pass_carry,
-        xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+    (Mc, M), _ = (
+        jax.lax.scan(
+            f=backward_pass,
+            init=backward_pass_carry,
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        )
+        if model.number_of_links() > 1
+        else [(Mc, M), None]
     )
 
     # Store the locked 6D rigid-body inertia matrix Mbb ∈ ℝ⁶ˣ⁶.

jaxsim/rbda/forward_kinematics.py

@@ -4,7 +4,7 @@ import jaxlie
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint, Quaternion
+from jaxsim.math import Adjoint
 
 from . import utils
 
@@ -42,7 +42,7 @@ def forward_kinematics_model(
 
     # Compute the base transform.
     W_H_B = jaxlie.SE3.from_rotation_and_translation(
-        rotation=jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(wxyz=W_Q_B)),
+        rotation=jaxlie.SO3(wxyz=W_Q_B),
         translation=W_p_B,
     )
 
@@ -75,10 +75,14 @@ def forward_kinematics_model(
 
         return (W_X_i,), None
 
-    (W_X_i,), _ = jax.lax.scan(
-        f=propagate_kinematics,
-        init=propagate_kinematics_carry,
-        xs=jnp.arange(start=1, stop=model.number_of_links()),
+    (W_X_i,), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_kinematics_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(W_X_i,), None]
     )
 
     return jax.vmap(Adjoint.to_transform)(W_X_i)