jaxsim 0.4.2.dev19__py3-none-any.whl → 0.4.2.dev31__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/api/contact.py

@@ -8,6 +8,7 @@ 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
@@ -411,3 +412,168 @@ 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(transform=jnp.eye(4))
+                O_Ẋ_W = W_Ẋ_W = jnp.zeros((6, 6))
+
+            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
+
+            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
+
+            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/api/model.py

@@ -576,6 +576,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,

jaxsim/api/references.py

@@ -30,6 +30,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 +38,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 +48,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         """
 
         return JaxSimModelReferences.build(
-            model=model, velocity_representation=velocity_representation
+            model=model, data=data, velocity_representation=velocity_representation
         )
 
     @staticmethod

{jaxsim-0.4.2.dev19.dist-info → jaxsim-0.4.2.dev31.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.2.dev19
+Version: 0.4.2.dev31
 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.dev19.dist-info → jaxsim-0.4.2.dev31.dist-info}/RECORD

@@ -1,21 +1,21 @@
 jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
-jaxsim/_version.py,sha256=fa--MfLuJkEPIj9KjXW-ToBHp9R1u3zovY8O3FGfsuY,426
+jaxsim/_version.py,sha256=NaaIf0Li8SvZ9MAX4mA2JwJaZ7w3LbcXaLdVp3HJJPg,426
 jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
 jaxsim/api/com.py,sha256=6TnYCvjmsJ2KLuw3NtZb0pay7ZwGKe9MKphYeQdjpQ0,13474
 jaxsim/api/common.py,sha256=Ubi6uAw3o6qbdU0TFGzUyHg98EnoMzrnlihrvrs95Sk,6653
-jaxsim/api/contact.py,sha256=EcOx_T94gZT3igtebmW9FJDpZYPEf-RwKfFN18JjOWM,13364
+jaxsim/api/contact.py,sha256=RdxPdwYbCpD_Wz0oiG9N569EoXtLiDGUAvCuoOJFYxc,19761
 jaxsim/api/data.py,sha256=T6m7-NteWrm-K3491Yo76xvlWtFCKqTzEE7Ughcasi8,27197
 jaxsim/api/frame.py,sha256=yQmhh8fckXnqzs7dQvojOzbuSanNGLwUTWUQDXbVtF4,12874
 jaxsim/api/joint.py,sha256=Pvg_It2iYA-jAQ2nOlFZxwmITiozO_f46G13BdQtHQ0,5106
 jaxsim/api/kin_dyn_parameters.py,sha256=CcfSg5Mc8qb1mZeMQ4AK_ffZIsK5yOl7tu397pFhcDA,29369
 jaxsim/api/link.py,sha256=hn7fbxaebHeXnvwEG9jZiWwzRcfdS8m-18LVsIG3S24,18479
-jaxsim/api/model.py,sha256=PMTSz00AIVopwiJ3zGBoYPTtkLH_beJCcQsX9wBE38I,61502
+jaxsim/api/model.py,sha256=8TTFVc7HuPLxIRq4h1pzXwuSlQ4sDVmL5KpiO-OnmnM,62662
 jaxsim/api/ode.py,sha256=NnLTBvpaT4kXnbjAghXIzLv9DTMJ8bele2iOlUQDv3Q,11028
 jaxsim/api/ode_data.py,sha256=9YZX-SK_KJtoIqG-zYWZsQInb2NA_LtxDn-jtLqm_3U,19759
-jaxsim/api/references.py,sha256=UA6kSQVBoq-bXSo99EOELf-_MD5MTy2zS0GtG3wQ410,16618
+jaxsim/api/references.py,sha256=e8M7o4DDgfxsuJzV0WMwGbgCD-kXTDKLfJre7iwciHQ,16821
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
 jaxsim/integrators/common.py,sha256=iwFykYZxdchqJcmcx8MFWEVijS5Hx9wCNKLKAJdF4gE,20103
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
@@ -61,8 +61,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=fLl1tY3DDb3lpIhG6BPqA5W34hM84oFzL-5cuz8k-68,11379
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=GOJQCJc5zwzoEGZB62wnWWGvUUQlXvDxz_A2Q-hFv7c,4027
-jaxsim-0.4.2.dev19.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.2.dev19.dist-info/METADATA,sha256=F_Rh5tV8N8VP-6NGTLxInNLuMTm1lRu9IRiu7r-jh7k,17250
-jaxsim-0.4.2.dev19.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
-jaxsim-0.4.2.dev19.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.2.dev19.dist-info/RECORD,,
+jaxsim-0.4.2.dev31.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.2.dev31.dist-info/METADATA,sha256=dnnmI39qsk5lb3hEZWe8mRZ1UntM_5KzSP5AL6S_jro,17250
+jaxsim-0.4.2.dev31.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+jaxsim-0.4.2.dev31.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.2.dev31.dist-info/RECORD,,