jaxsim 0.2.1.dev115__py3-none-any.whl → 0.2.1.dev123__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/link.py

@@ -3,6 +3,7 @@ from typing import Sequence
 
 import jax
 import jax.numpy as jnp
+import jax.scipy.linalg
 import jaxlie
 import numpy as np
 
@@ -337,6 +338,187 @@ def velocity(
     return O_J_WL_I @ I_ν
 
 
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def jacobian_derivative(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    """
+    Compute the derivative of the free-floating jacobian of the link.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_index: The index of the link.
+        output_vel_repr:
+            The output velocity representation of the free-floating jacobian derivative.
+
+    Returns:
+        The derivative of the 6×(6+n) free-floating jacobian of the link.
+
+    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
+    )
+
+    # 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(),
+    )
+
+    # 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[link_index]
+    B_J̇_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J̇_full_WX_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:
+
+            W_H_B = data.base_transform()
+            B_X_W = jaxsim.math.Adjoint.from_transform(transform=W_H_B, inverse=True)
+
+            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)
+
+            # 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 = 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)
+
+            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)
+
+            # 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:
+
+            W_H_B = data.base_transform()
+            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)
+
+        case VelRepr.Body:
+
+            O_X_B = L_X_B = jaxsim.math.Adjoint.from_transform(
+                transform=B_H_L[link_index, :, :], 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 = js.link.velocity(model=model, data=data, link_index=link_index)
+
+            O_Ẋ_B = L_Ẋ_B = -L_X_B @ jaxsim.math.Cross.vx(B_X_L @ L_v_WL - B_v_WB)
+
+        case VelRepr.Mixed:
+
+            W_H_B = data.base_transform()
+            W_H_L = W_H_B @ B_H_L[link_index, :, :]
+            LW_H_L = W_H_L.at[0:3, 3].set(jnp.zeros(3))
+            LW_H_B = LW_H_L @ jaxsim.math.Transform.inverse(B_H_L[link_index, :, :])
+
+            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):
+                LW_v_WL = js.link.velocity(
+                    model=model, data=data, link_index=link_index
+                )
+                LW_v_W_LW = LW_v_WL.at[3:6].set(jnp.zeros(3))
+
+            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(B_X_LW @ LW_v_B_LW)
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+    # =============================================================
+    # Express the Jacobian derivative in the target representations
+    # =============================================================
+
+    # The derivative of the equation to change the input and output representations
+    # of the Jacobian derivative needs the computation of the plain link Jacobian.
+    # Compute here the full Jacobian of the model...
+    B_J_full_WL_B, _ = jaxsim.rbda.jacobian_full_doubly_left(
+        model=model,
+        joint_positions=data.joint_positions(),
+    )
+
+    # ... and extract the link Jacobian using the boolean support body array.
+    B_J_WL_B = jnp.hstack([jnp.ones(5), κb]) * B_J_full_WL_B
+
+    # Sum all the components that form the Jacobian derivative in the target
+    # input/output velocity representations.
+    O_J̇_WL_I = jnp.zeros(shape=(6, 6 + model.dofs()))
+    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
+
+
 @jax.jit
 def bias_acceleration(
     model: js.model.JaxSimModel,

jaxsim/rbda/__init__.py

@@ -2,6 +2,10 @@ from .aba import aba
 from .collidable_points import collidable_points_pos_vel
 from .crba import crba
 from .forward_kinematics import forward_kinematics, forward_kinematics_model
-from .jacobian import jacobian, jacobian_full_doubly_left
+from .jacobian import (
+    jacobian,
+    jacobian_derivative_full_doubly_left,
+    jacobian_full_doubly_left,
+)
 from .rnea import rnea
 from .soft_contacts import SoftContacts, SoftContactsParams

jaxsim/rbda/jacobian.py

@@ -4,7 +4,7 @@ import numpy as np
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint
+from jaxsim.math import Adjoint, Cross
 
 from . import utils
 
@@ -199,3 +199,120 @@ def jacobian_full_doubly_left(
     B_J_full_WL_B = J.squeeze().astype(float)
 
     return B_J_full_WL_B, B_H_L
+
+
+def jacobian_derivative_full_doubly_left(
+    model: js.model.JaxSimModel,
+    *,
+    joint_positions: jtp.VectorLike,
+    joint_velocities: jtp.VectorLike,
+) -> tuple[jtp.Matrix, jtp.Array]:
+    r"""
+    Compute the derivative of the doubly-left full free-floating Jacobian of a model.
+
+    The derivative of the full Jacobian is a 6x(6+n) matrix with all the columns filled.
+    It is useful to run the algorithm once, and then extract the link Jacobian
+    derivative by filtering the columns of the full Jacobian using the support
+    parent array :math:`\kappa(i)` of the link.
+
+    Args:
+        model: The model to consider.
+        joint_positions: The positions of the joints.
+        joint_velocities: The velocities of the joints.
+
+    Returns:
+        The derivative of the doubly-left full free-floating Jacobian of a model.
+    """
+
+    _, _, s, _, ṡ, _, _, _, _, _ = utils.process_inputs(
+        model=model, joint_positions=joint_positions, joint_velocities=joint_velocities
+    )
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # 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=s, base_transform=jnp.eye(4)
+    )
+
+    # Allocate the buffer of 6D transform base -> link.
+    B_X_i = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    B_X_i = B_X_i.at[0].set(jnp.eye(6))
+
+    # Allocate the buffer of 6D transform derivatives base -> link.
+    B_Ẋ_i = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+
+    # Allocate the buffer of the 6D link velocity in body-fixed representation.
+    B_v_Bi = jnp.zeros(shape=(model.number_of_links(), 6))
+
+    # Helper to compute the time derivative of the adjoint matrix.
+    def A_Ẋ_B(A_X_B: jtp.Matrix, B_v_AB: jtp.Vector) -> jtp.Matrix:
+        return A_X_B @ Cross.vx(B_v_AB).squeeze()
+
+    # ============================================
+    # Compute doubly-left full Jacobian derivative
+    # ============================================
+
+    # Allocate the Jacobian matrix.
+    J̇ = jnp.zeros(shape=(6, 6 + model.dofs()))
+
+    ComputeFullJacobianDerivativeCarry = tuple[
+        jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
+    ]
+
+    compute_full_jacobian_derivative_carry: ComputeFullJacobianDerivativeCarry = (
+        B_v_Bi,
+        B_X_i,
+        B_Ẋ_i,
+        J̇,
+    )
+
+    def compute_full_jacobian_derivative(
+        carry: ComputeFullJacobianDerivativeCarry, i: jtp.Int
+    ) -> tuple[ComputeFullJacobianDerivativeCarry, None]:
+
+        ii = i - 1
+        B_v_Bi, B_X_i, B_Ẋ_i, J̇ = carry
+
+        # Compute the base (0) to link (i) adjoint matrix.
+        B_Xi_i = B_X_i[λ[i]] @ Adjoint.inverse(i_X_λi[i])
+        B_X_i = B_X_i.at[i].set(B_Xi_i)
+
+        # Compute the body-fixed velocity of the link.
+        B_vi_Bi = B_v_Bi[λ[i]] + B_X_i[i] @ S[i].squeeze() * ṡ[ii]
+        B_v_Bi = B_v_Bi.at[i].set(B_vi_Bi)
+
+        # Compute the base (0) to link (i) adjoint matrix derivative.
+        i_Xi_B = Adjoint.inverse(B_Xi_i)
+        B_Ẋi_i = A_Ẋ_B(A_X_B=B_Xi_i, B_v_AB=i_Xi_B @ B_vi_Bi)
+        B_Ẋ_i = B_Ẋ_i.at[i].set(B_Ẋi_i)
+
+        # Compute the ii-th column of the B_Ṡ_BL(s) matrix.
+        B_Ṡii_BL = B_Ẋ_i[i] @ S[i]
+        J̇ = J̇.at[0:6, 6 + ii].set(B_Ṡii_BL.squeeze())
+
+        return (B_v_Bi, B_X_i, B_Ẋ_i, J̇), None
+
+    (_, B_X_i, B_Ẋ_i, J̇), _ = (
+        jax.lax.scan(
+            f=compute_full_jacobian_derivative,
+            init=compute_full_jacobian_derivative_carry,
+            xs=np.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(_, B_X_i, B_Ẋ_i, J̇), None]
+    )
+
+    # Convert adjoints to SE(3) transforms.
+    # Returning them here prevents calling FK in case the output representation
+    # of the Jacobian needs to be changed.
+    B_H_L = jax.vmap(lambda B_X_L: Adjoint.to_transform(B_X_L))(B_X_i)
+
+    # Adjust shape of doubly-left free-floating full Jacobian derivative.
+    B_J̇_full_WL_B = J̇.squeeze().astype(float)
+
+    return B_J̇_full_WL_B, B_H_L

{jaxsim-0.2.1.dev115.dist-info → jaxsim-0.2.1.dev123.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev115
+Version: 0.2.1.dev123
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.2.1.dev115.dist-info → jaxsim-0.2.1.dev123.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=EYfLqc9WPDAaN-SdBrvQC_6Ng8tlzXg8pbUvy5VJpB4,428
+jaxsim/_version.py,sha256=jAmAP2ZYBVoFw5W8EyAYOEpIKED5Jy8XK259JnSs0Fk,428
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
@@ -10,7 +10,7 @@ jaxsim/api/data.py,sha256=xfKJz6Rw0YTk-EHCGiT8BFQrs_ggOz01lRi1Qh1mb28,27256
 jaxsim/api/frame.py,sha256=0YXOrGmx3cSQqa4_Ky-n6zyup3I3xvXNEgub-Bc5xUw,6222
 jaxsim/api/joint.py,sha256=-5DogPg4g4mmLckyVIVNjwv-Rxz0IWS7_md9nDlhPWA,4581
 jaxsim/api/kin_dyn_parameters.py,sha256=zMca7OmCsCWK_cavLTSZSeYh9Qu1-409cdsyWvWPAUQ,26090
-jaxsim/api/link.py,sha256=rypTwkMf9HJ5UuAtHRJh0LqqdJWcLKTtTjWcjduEsF0,9842
+jaxsim/api/link.py,sha256=oW5-DShmmeCRk3JOJIwzo3HbWuNGmpm_wBJ4fkmrROM,16645
 jaxsim/api/model.py,sha256=1HlQ5FMzeJAk-cE1pmELgVjzMYUX9-iipw3N4WssAL4,55435
 jaxsim/api/ode.py,sha256=BfvV_14uu0szWecoDiV8rTu-dvSFLK7eyrO38ZqHB_w,10157
 jaxsim/api/ode_data.py,sha256=D6FzMkvY_qNuoFEImyp7sxAk-0pJOd3oZeSr9bBTcLk,23089
@@ -43,12 +43,12 @@ jaxsim/parsers/descriptions/model.py,sha256=uO5xOJtViihVPnSSsmfQJvCh45ANyi9KYAzL
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
 jaxsim/parsers/rod/parser.py,sha256=Q13TOkmpU0SHpgSV8WRYWb290aPNNLsaz4eMlD4Mq5w,13525
 jaxsim/parsers/rod/utils.py,sha256=9oO4YsQRaR2v700IkNOXRnPpn5i4N8HFfgjPkMLK2mc,5732
-jaxsim/rbda/__init__.py,sha256=HLwxeU-IxaRpFGUCSQv-LDv20JHTt3Xj7ELiRbRieS8,319
+jaxsim/rbda/__init__.py,sha256=MqEZwzu8SHPAlIFHmSXmCjehuOJGRX58OrBVAbBVMwg,374
 jaxsim/rbda/aba.py,sha256=0OoCzHhf1v-qqr1y5PIrD7_mPwAlid0fjXxUrIa5E_s,9118
 jaxsim/rbda/collidable_points.py,sha256=4ZNJbEj2nEi15jBLR-GNbdaqKgkN58FBgqd_TXupEgg,4948
 jaxsim/rbda/crba.py,sha256=GodskOZjtrSlbQAqxRv1un_706O7BaJK-U2qa18vJk8,4741
 jaxsim/rbda/forward_kinematics.py,sha256=OHugNU7C0UxYAW0o1rqH1ZgniSwurz6L1T1MJxfxq08,3418
-jaxsim/rbda/jacobian.py,sha256=9LGGy9ya5m5U0mBmV1NFH5XYZpEMYbx74qnYBvZs7Ok,6360
+jaxsim/rbda/jacobian.py,sha256=M79bGir-2w_iJ2GurYhOGgMfJnp7ZMOCW6AeeWKK8iM,10745
 jaxsim/rbda/rnea.py,sha256=DjwkvXQVUSUclM3Uy3UPZ2tao91R5dGd4o7TsS2qObI,7650
 jaxsim/rbda/soft_contacts.py,sha256=52zJOF31hFpqoaOednTvi8j_UxhRcdGNjzOPb2v2MPc,11257
 jaxsim/rbda/utils.py,sha256=zpbFM2Iq8cntku0BFVu9nfEqZhInCWi9D2INT6MFEI8,5003
@@ -58,8 +58,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=EJMcblYKUjxw9HJShVf81Ig3pHUJno6Dx6h-RnY--wM,2040
-jaxsim-0.2.1.dev115.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.2.1.dev115.dist-info/METADATA,sha256=hrMw2YCSDcAoA7rLTnNYKwmGmHPDOuaohkwaATmW6u0,9745
-jaxsim-0.2.1.dev115.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.1.dev115.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.1.dev115.dist-info/RECORD,,
+jaxsim-0.2.1.dev123.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.2.1.dev123.dist-info/METADATA,sha256=hOUnUjhHwGvdZVMR_-z5pKFsd681swunNdYkqgC338I,9745
+jaxsim-0.2.1.dev123.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.1.dev123.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.1.dev123.dist-info/RECORD,,