jaxsim 0.4.2__tar.gz → 0.4.2.dev14__tar.gz

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.2
+Version: 0.4.2.dev14
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/examples/PD_controller.ipynb

@@ -37,9 +37,16 @@
     "\n",
     "    # Install dependencies for visualization on Colab and ReadTheDocs\n",
     "    !sudo apt update\n",
-    "    !apt install libosmesa6-dev\n",
+    "    !apt install xvfb ffmpeg python-opengl\n",
+    "    !{sys.executable} -m pip install pyvirtualdisplay\n",
     "    clear_output()\n",
     "\n",
+    "    # Start Xvfb and set \"DISPLAY\" environment variable\n",
+    "    import pyvirtualdisplay\n",
+    "\n",
+    "    d = pyvirtualdisplay.Display()\n",
+    "    d.start()\n",
+    "\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
@@ -147,21 +154,51 @@
    "outputs": [],
    "source": [
     "# @title Set up MuJoCo renderer\n",
+    "import subprocess\n",
+    "\n",
     "os.environ[\"MUJOCO_GL\"] = \"osmesa\"\n",
     "\n",
+    "if IS_COLAB:\n",
+    "    if subprocess.run(\"nvidia-smi\").returncode:\n",
+    "        raise RuntimeError(\n",
+    "            \"Cannot communicate with GPU. \"\n",
+    "            \"Make sure you are using a GPU Colab runtime. \"\n",
+    "            \"Go to the Runtime menu and select Choose runtime type.\"\n",
+    "        )\n",
+    "\n",
+    "    # Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
+    "    # This is usually installed as part of an Nvidia driver package, but the Colab\n",
+    "    # kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
+    "    # (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
+    "    NVIDIA_ICD_CONFIG_PATH = \"/usr/share/glvnd/egl_vendor.d/10_nvidia.json\"\n",
+    "    if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
+    "        with open(NVIDIA_ICD_CONFIG_PATH, \"w\") as f:\n",
+    "            f.write(\n",
+    "                \"\"\"{\n",
+    "      \"file_format_version\" : \"1.0.0\",\n",
+    "      \"ICD\" : {\n",
+    "         \"library_path\" : \"libEGL_nvidia.so.0\"\n",
+    "     }\n",
+    "  }\n",
+    "  \"\"\"\n",
+    "            )\n",
+    "\n",
+    "    os.environ[\"MUJOCO_GL\"] = \"egl\"\n",
+    "\n",
+    "\n",
     "from jaxsim.mujoco import MujocoModelHelper, MujocoVideoRecorder\n",
-    "from jaxsim.mujoco.loaders import UrdfToMjcf, MujocoCamera\n",
-    "\n",
-    "mjcf_string, assets = UrdfToMjcf.convert(\n",
-    "    urdf=model.built_from,\n",
-    "    cameras=MujocoCamera.build_from_target_view(\n",
-    "        camera_name=\"cartpole_camera\",\n",
-    "        lookat=jnp.array([0.0, data.joint_positions()[0], 1.2]),\n",
-    "        distance=3,\n",
-    "        azimut=150,\n",
-    "        elevation=-10,\n",
-    "    ),\n",
-    ")\n",
+    "from jaxsim.mujoco.loaders import UrdfToMjcf\n",
+    "\n",
+    "camera = {\n",
+    "    \"name\": \"cartpole_camera\",\n",
+    "    \"mode\": \"fixed\",\n",
+    "    \"pos\": \"3.954 3.533 2.343\",\n",
+    "    \"xyaxes\": \"-0.594 0.804 -0.000 -0.163 -0.120 0.979\",\n",
+    "    \"fovy\": \"60\",\n",
+    "}\n",
+    "\n",
+    "mjcf_string, assets = UrdfToMjcf.convert(urdf=model.built_from, cameras=camera)\n",
+    "\n",
     "mj_model_helper = MujocoModelHelper.build_from_xml(\n",
     "    mjcf_description=mjcf_string, assets=assets\n",
     ")\n",

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/pyproject.toml

@@ -156,6 +156,7 @@ ignore = [
     "E501", # Line too long
     "E731", # Do not assign a `lambda` expression, use a `def`
     "E741", # Ambiguous variable name
+    "F841", # Local variable is assigned to but never used
     "I001", # Import block is unsorted or unformatted
     "RUF003", # Ambigous unicode character in comment
 ]

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.2'
-__version_tuple__ = version_tuple = (0, 4, 2)
+__version__ = version = '0.4.2.dev14'
+__version_tuple__ = version_tuple = (0, 4, 2, 'dev14')

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/src/jaxsim/api/com.py

@@ -137,9 +137,9 @@ def centroidal_momentum_jacobian(
 
     match data.velocity_representation:
         case VelRepr.Inertial | VelRepr.Mixed:
-            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)  # noqa: F841
+            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
         case VelRepr.Body:
-            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)  # noqa: F841
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
         case _:
             raise ValueError(data.velocity_representation)
 
@@ -172,9 +172,9 @@ def locked_centroidal_spatial_inertia(
 
     match data.velocity_representation:
         case VelRepr.Inertial | VelRepr.Mixed:
-            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)  # noqa: F841
+            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
         case VelRepr.Body:
-            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)  # noqa: F841
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
         case _:
             raise ValueError(data.velocity_representation)
 
@@ -290,14 +290,14 @@ def bias_acceleration(
 
         case VelRepr.Inertial:
 
-            C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL  # noqa: F841
-            C_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
+            C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL
+            C_v_WC = W_v_WW = jnp.zeros(6)
 
-            L_H_C = L_H_W = jax.vmap(  # noqa: F841
+            L_H_C = L_H_W = jax.vmap(
                 lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
             )(W_H_L)
 
-            L_v_LC = L_v_LW = jax.vmap(  # noqa: F841
+            L_v_LC = L_v_LW = jax.vmap(
                 lambda i: -js.link.velocity(
                     model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
                 )
@@ -314,9 +314,9 @@ def bias_acceleration(
 
         case VelRepr.Mixed:
 
-            C_v̇_WL = LW_v̇_bias_WL = v̇_bias_WL  # noqa: F841
+            C_v̇_WL = LW_v̇_bias_WL = v̇_bias_WL
 
-            C_v_WC = LW_v_W_LW = jax.vmap(  # noqa: F841
+            C_v_WC = LW_v_W_LW = jax.vmap(
                 lambda i: js.link.velocity(
                     model=model, data=data, link_index=i, output_vel_repr=VelRepr.Mixed
                 )
@@ -324,13 +324,13 @@ def bias_acceleration(
                 .set(jnp.zeros(3))
             )(jnp.arange(model.number_of_links()))
 
-            L_H_C = L_H_LW = jax.vmap(  # noqa: F841
+            L_H_C = L_H_LW = jax.vmap(
                 lambda W_H_L: jaxsim.math.Transform.inverse(
                     W_H_L.at[0:3, 3].set(jnp.zeros(3))
                 )
             )(W_H_L)
 
-            L_v_LC = L_v_L_LW = jax.vmap(  # noqa: F841
+            L_v_LC = L_v_L_LW = jax.vmap(
                 lambda i: -js.link.velocity(
                     model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
                 )

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/src/jaxsim/api/contact.py

@@ -8,7 +8,6 @@ import jax.numpy as jnp
 import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint, Cross, Transform
 from jaxsim.rbda.contacts.soft import SoftContactsParams
 
 from .common import VelRepr
@@ -412,170 +411,3 @@ def jacobian(
             raise ValueError(output_vel_repr)
 
     return O_J_WC
-
-
-@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
-def jacobian_derivative(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
-    *,
-    output_vel_repr: VelRepr | None = None,
-) -> jtp.Matrix:
-    r"""
-    Compute the derivative of the free-floating jacobian of the contact points.
-
-    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 derivative.
-
-    Returns:
-        The derivative of the :math:`6 \times (6+n)` free-floating jacobian of the contact points.
-
-    Note:
-        The input representation of the free-floating jacobian derivative is the active
-        velocity representation.
-    """
-
-    output_vel_repr = (
-        output_vel_repr if output_vel_repr is not None else data.velocity_representation
-    )
-
-    # Get the index of the parent link and the position of the collidable point.
-    parent_link_idxs = jnp.array(model.kin_dyn_parameters.contact_parameters.body)
-    L_p_Ci = jnp.array(model.kin_dyn_parameters.contact_parameters.point)
-    contact_idxs = jnp.arange(L_p_Ci.shape[0])
-
-    # Get the transforms of all the parent links.
-    W_H_Li = js.model.forward_kinematics(model=model, data=data)
-
-    # =====================================================
-    # Compute quantities to adjust the input representation
-    # =====================================================
-
-    def compute_T(model: js.model.JaxSimModel, X: jtp.Matrix) -> jtp.Matrix:
-        In = jnp.eye(model.dofs())
-        T = jax.scipy.linalg.block_diag(X, In)
-        return T
-
-    def compute_Ṫ(model: js.model.JaxSimModel, Ẋ: jtp.Matrix) -> jtp.Matrix:
-        On = jnp.zeros(shape=(model.dofs(), model.dofs()))
-        Ṫ = jax.scipy.linalg.block_diag(Ẋ, On)
-        return Ṫ
-
-    # Compute the operator to change the representation of ν, and its
-    # time derivative.
-    match data.velocity_representation:
-        case VelRepr.Inertial:
-            W_H_W = jnp.eye(4)
-            W_X_W = Adjoint.from_transform(transform=W_H_W)
-            W_Ẋ_W = jnp.zeros((6, 6))
-
-            T = compute_T(model=model, X=W_X_W)
-            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_W)
-
-        case VelRepr.Body:
-            W_H_B = data.base_transform()
-            W_X_B = Adjoint.from_transform(transform=W_H_B)
-            B_v_WB = data.base_velocity()
-            B_vx_WB = Cross.vx(B_v_WB)
-            W_Ẋ_B = W_X_B @ B_vx_WB
-
-            T = compute_T(model=model, X=W_X_B)
-            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_B)
-
-        case VelRepr.Mixed:
-            W_H_B = data.base_transform()
-            W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
-            W_X_BW = Adjoint.from_transform(transform=W_H_BW)
-            BW_v_WB = data.base_velocity()
-            BW_v_W_BW = BW_v_WB.at[3:6].set(jnp.zeros(3))
-            BW_vx_W_BW = Cross.vx(BW_v_W_BW)
-            W_Ẋ_BW = W_X_BW @ BW_vx_W_BW
-
-            T = compute_T(model=model, X=W_X_BW)
-            Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_BW)
-
-        case _:
-            raise ValueError(data.velocity_representation)
-
-    # =====================================================
-    # Compute quantities to adjust the output representation
-    # =====================================================
-
-    with data.switch_velocity_representation(VelRepr.Inertial):
-        # Compute the Jacobian of the parent link in inertial representation.
-        W_J_WL_W = js.model.generalized_free_floating_jacobian(
-            model=model,
-            data=data,
-            output_vel_repr=VelRepr.Inertial,
-        )
-        # Compute the Jacobian derivative of the parent link in inertial representation.
-        W_J̇_WL_W = js.model.generalized_free_floating_jacobian_derivative(
-            model=model,
-            data=data,
-            output_vel_repr=VelRepr.Inertial,
-        )
-
-    # Get the Jacobian of the collidable points in the mixed representation.
-    with data.switch_velocity_representation(VelRepr.Mixed):
-        CW_J_WC_BW = jacobian(
-            model=model,
-            data=data,
-            output_vel_repr=VelRepr.Mixed,
-        )
-
-    def compute_O_J̇_WC_I(
-        L_p_C: jtp.Vector,
-        contact_idx: jtp.Int,
-        CW_J_WC_BW: jtp.Matrix,
-        W_H_L: jtp.Matrix,
-    ) -> jtp.Matrix:
-
-        parent_link_idx = parent_link_idxs[contact_idx]
-
-        match output_vel_repr:
-            case VelRepr.Inertial:
-                O_X_W = W_X_W = Adjoint.from_transform(  # noqa: F841
-                    transform=jnp.eye(4)
-                )
-                O_Ẋ_W = W_Ẋ_W = jnp.zeros((6, 6))  # noqa: F841
-
-            case VelRepr.Body:
-                L_H_C = Transform.from_rotation_and_translation(translation=L_p_C)
-                W_H_C = W_H_L[parent_link_idx] @ L_H_C
-                O_X_W = C_X_W = Adjoint.from_transform(transform=W_H_C, inverse=True)
-                with data.switch_velocity_representation(VelRepr.Inertial):
-                    W_nu = data.generalized_velocity()
-                W_v_WC = W_J_WL_W[parent_link_idx] @ W_nu
-                W_vx_WC = Cross.vx(W_v_WC)
-                O_Ẋ_W = C_Ẋ_W = -C_X_W @ W_vx_WC  # noqa: F841
-
-            case VelRepr.Mixed:
-                L_H_C = Transform.from_rotation_and_translation(translation=L_p_C)
-                W_H_C = W_H_L[parent_link_idx] @ L_H_C
-                W_H_CW = W_H_C.at[0:3, 0:3].set(jnp.eye(3))
-                CW_H_W = Transform.inverse(W_H_CW)
-                O_X_W = CW_X_W = Adjoint.from_transform(transform=CW_H_W)
-                with data.switch_velocity_representation(VelRepr.Mixed):
-                    CW_v_WC = CW_J_WC_BW @ data.generalized_velocity()
-                W_v_W_CW = jnp.zeros(6).at[0:3].set(CW_v_WC[0:3])
-                W_vx_W_CW = Cross.vx(W_v_W_CW)
-                O_Ẋ_W = CW_Ẋ_W = -CW_X_W @ W_vx_W_CW  # noqa: F841
-
-            case _:
-                raise ValueError(output_vel_repr)
-
-        O_J̇_WC_I = jnp.zeros(shape=(6, 6 + model.dofs()))
-        O_J̇_WC_I += O_Ẋ_W @ W_J_WL_W[parent_link_idx] @ T
-        O_J̇_WC_I += O_X_W @ W_J̇_WL_W[parent_link_idx] @ T
-        O_J̇_WC_I += O_X_W @ W_J_WL_W[parent_link_idx] @ Ṫ
-
-        return O_J̇_WC_I
-
-    O_J̇_WC = jax.vmap(compute_O_J̇_WC_I, in_axes=(0, 0, 0, None))(
-        L_p_Ci, contact_idxs, CW_J_WC_BW, W_H_Li
-    )
-
-    return O_J̇_WC

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/src/jaxsim/api/data.py

@@ -12,6 +12,7 @@ import jaxlie
 import jaxsim.api as js
 import jaxsim.rbda
 import jaxsim.typing as jtp
+from jaxsim.math import Quaternion
 from jaxsim.rbda.contacts.soft import SoftContacts
 from jaxsim.utils import Mutability
 from jaxsim.utils.tracing import not_tracing
@@ -190,7 +191,9 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         W_H_B = jaxlie.SE3.from_rotation_and_translation(
             translation=base_position,
-            rotation=jaxlie.SO3(wxyz=base_quaternion),
+            rotation=jaxlie.SO3.from_quaternion_xyzw(
+                base_quaternion[jnp.array([1, 2, 3, 0])]
+            ),
         ).as_matrix()
 
         v_WB = JaxSimModelData.other_representation_to_inertial(
@@ -377,7 +380,13 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             on_false=W_Q_B / jnp.linalg.norm(W_Q_B),
         )
 
-        return (W_Q_B if not dcm else jaxlie.SO3(wxyz=W_Q_B).as_matrix()).astype(float)
+        return (
+            W_Q_B
+            if not dcm
+            else jaxlie.SO3.from_quaternion_xyzw(
+                Quaternion.to_xyzw(wxyz=W_Q_B)
+            ).as_matrix()
+        ).astype(float)
 
     @jax.jit
     def base_transform(self) -> jtp.Matrix:

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/src/jaxsim/api/frame.py

@@ -384,7 +384,7 @@ def jacobian_derivative(
                 W_nu = data.generalized_velocity()
             W_v_WF = W_J_WL_W @ W_nu
             W_vx_WF = Cross.vx(W_v_WF)
-            O_Ẋ_W = F_Ẋ_W = -F_X_W @ W_vx_WF  # noqa: F841
+            O_Ẋ_W = F_Ẋ_W = -F_X_W @ W_vx_WF
 
         case VelRepr.Mixed:
             W_H_F = transform(model=model, data=data, frame_index=frame_index)
@@ -401,7 +401,7 @@ def jacobian_derivative(
                 FW_v_WF = FW_J_WF_FW @ data.generalized_velocity()
             W_v_W_FW = jnp.zeros(6).at[0:3].set(FW_v_WF[0:3])
             W_vx_W_FW = Cross.vx(W_v_W_FW)
-            O_Ẋ_W = FW_Ẋ_W = -FW_X_W @ W_vx_W_FW  # noqa: F841
+            O_Ẋ_W = FW_Ẋ_W = -FW_X_W @ W_vx_W_FW
 
         case _:
             raise ValueError(output_vel_repr)

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/src/jaxsim/api/link.py

@@ -288,7 +288,7 @@ def jacobian(
         case VelRepr.Inertial:
             W_H_B = data.base_transform()
             B_X_W = Adjoint.from_transform(transform=W_H_B, inverse=True)
-            B_J_WL_I = B_J_WL_W = B_J_WL_B @ jax.scipy.linalg.block_diag(  # noqa: F841
+            B_J_WL_I = B_J_WL_W = B_J_WL_B @ jax.scipy.linalg.block_diag(
                 B_X_W, jnp.eye(model.dofs())
             )
 
@@ -299,7 +299,7 @@ def jacobian(
             W_R_B = data.base_orientation(dcm=True)
             BW_H_B = jnp.eye(4).at[0:3, 0:3].set(W_R_B)
             B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
-            B_J_WL_I = B_J_WL_BW = B_J_WL_B @ jax.scipy.linalg.block_diag(  # noqa: F841
+            B_J_WL_I = B_J_WL_BW = B_J_WL_B @ jax.scipy.linalg.block_diag(
                 B_X_BW, jnp.eye(model.dofs())
             )
 
@@ -313,7 +313,7 @@ def jacobian(
         case VelRepr.Inertial:
             W_H_B = data.base_transform()
             W_X_B = Adjoint.from_transform(transform=W_H_B)
-            O_J_WL_I = W_J_WL_I = W_X_B @ B_J_WL_I  # noqa: F841
+            O_J_WL_I = W_J_WL_I = W_X_B @ B_J_WL_I
 
         case VelRepr.Body:
             L_X_B = Adjoint.from_transform(transform=B_H_L, inverse=True)
@@ -505,7 +505,7 @@ def jacobian_derivative(
             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
+            O_Ẋ_B = W_Ẋ_B = W_X_B @ jaxsim.math.Cross.vx(B_v_WB)
 
         case VelRepr.Body:
 
@@ -519,9 +519,7 @@ def jacobian_derivative(
                 B_v_WB = data.base_velocity()
                 L_v_WL = js.link.velocity(model=model, data=data, link_index=link_index)
 
-            O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
-                B_X_L @ L_v_WL - B_v_WB
-            )
+            O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(B_X_L @ L_v_WL - B_v_WB)
 
         case VelRepr.Mixed:
 
@@ -546,9 +544,8 @@ def jacobian_derivative(
             LW_v_LW_L = LW_v_WL - LW_v_W_LW
             LW_v_B_LW = LW_v_WL - LW_X_B @ B_v_WB - LW_v_LW_L
 
-            O_Ẋ_B = LW_Ẋ_B = -LW_X_B @ jaxsim.math.Cross.vx(  # noqa: F841
-                B_X_LW @ LW_v_B_LW
-            )
+            O_Ẋ_B = LW_Ẋ_B = -LW_X_B @ jaxsim.math.Cross.vx(B_X_LW @ LW_v_B_LW)
+
         case _:
             raise ValueError(output_vel_repr)
 

{jaxsim-0.4.2 → jaxsim-0.4.2.dev14}/src/jaxsim/api/model.py

@@ -495,9 +495,8 @@ 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 = (  # noqa: F841
-                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 = B_J_full_WX_B @ jax.scipy.linalg.block_diag(
+                B_X_W, jnp.eye(model.dofs())
             )
 
         case VelRepr.Body:
@@ -510,7 +509,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 = (  # noqa: F841
+            B_J_full_WX_I = B_J_full_WX_BW = (
                 B_J_full_WX_B
                 @ jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
             )
@@ -543,13 +542,11 @@ 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(  # noqa: F841
-                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(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(  # noqa: F841
+            O_J_WL_I = L_J_WL_I = jax.vmap(
                 lambda B_H_L, B_J_WL_I: jaxsim.math.Adjoint.from_transform(
                     B_H_L, inverse=True
                 )
@@ -568,7 +565,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(  # noqa: F841
+            O_J_WL_I = LW_J_WL_I = jax.vmap(
                 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)
@@ -579,41 +576,6 @@ 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,
@@ -759,8 +721,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)  # noqa: F841
-            W_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
+            W_H_C = W_H_W = jnp.eye(4)
+            W_v_WC = W_v_WW = jnp.zeros(6)
 
         case VelRepr.Body:
             # In this case C=B
@@ -770,9 +732,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))  # noqa: F841
+            W_H_C = W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
             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
+            W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)
 
         case _:
             raise ValueError(data.velocity_representation)
@@ -1127,8 +1089,8 @@ def inverse_dynamics(
 
     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
+            W_H_C = W_H_W = jnp.eye(4)
+            W_v_WC = W_v_WW = jnp.zeros(6)
 
         case VelRepr.Body:
             W_H_C = W_H_B = data.base_transform()
@@ -1137,9 +1099,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))  # noqa: F841
+            W_H_C = W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
             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
+            W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)
 
         case _:
             raise ValueError(data.velocity_representation)
@@ -1574,15 +1536,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)  # noqa: F841
-            W_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
+            W_H_C = W_H_W = jnp.eye(4)
+            W_v_WC = W_v_WW = jnp.zeros(6)
             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
+                W_v_WC = W_v_WB = data.base_velocity()
             with data.switch_velocity_representation(VelRepr.Body):
                 C_v_WB = B_v_WB = data.base_velocity()
 
@@ -1593,9 +1555,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  # noqa: F841
+                W_v_WC = W_v_W_BW = W_X_BW @ BW_v_W_BW
             with data.switch_velocity_representation(VelRepr.Mixed):
-                C_v_WB = BW_v_WB = data.base_velocity()  # noqa: F841
+                C_v_WB = BW_v_WB = data.base_velocity()
 
         case _:
             raise ValueError(data.velocity_representation)
@@ -1703,12 +1665,8 @@ def link_bias_accelerations(
 
     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)
-            )
+            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))
 
         case VelRepr.Inertial:
             C_H_L = W_H_L = js.model.forward_kinematics(model=model, data=data)
@@ -1718,9 +1676,7 @@ 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(  # noqa: F841
-                lambda v: v.at[0:3].set(jnp.zeros(3))
-            )(L_v_WL)
+            L_v_CL = L_v_LW_L = jax.vmap(lambda v: v.at[0:3].set(jnp.zeros(3)))(L_v_WL)
 
         case _:
             raise ValueError(data.velocity_representation)