jaxsim 0.4.1.dev26__tar.gz → 0.4.2.dev12__tar.gz

{jaxsim-0.4.1.dev26 → jaxsim-0.4.2.dev12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.1.dev26
+Version: 0.4.2.dev12
 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.1.dev26 → jaxsim-0.4.2.dev12}/src/jaxsim/_version.py

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

{jaxsim-0.4.1.dev26 → jaxsim-0.4.2.dev12}/src/jaxsim/api/frame.py

@@ -7,7 +7,7 @@ import jax.numpy as jnp
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim import exceptions
-from jaxsim.math import Adjoint, Transform
+from jaxsim.math import Adjoint, Cross, Transform
 
 from .common import VelRepr
 
@@ -257,3 +257,158 @@ def jacobian(
             raise ValueError(output_vel_repr)
 
     return O_J_WL_I
+
+
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def jacobian_derivative(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Matrix:
+    r"""
+    Compute the derivative of the free-floating jacobian of the frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame.
+        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 frame.
+
+    Note:
+        The input representation of the free-floating jacobian derivative is the active
+        velocity representation.
+    """
+
+    n_l = model.number_of_links()
+    n_f = len(model.frame_names())
+
+    exceptions.raise_value_error_if(
+        condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
+        msg="Invalid frame index '{idx}'",
+        idx=frame_index,
+    )
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Get the index of the parent link.
+    L = idx_of_parent_link(model=model, frame_index=frame_index)
+
+    with data.switch_velocity_representation(VelRepr.Inertial):
+        # Compute the Jacobian of the parent link in inertial representation.
+        W_J_WL_W = js.link.jacobian(
+            model=model,
+            data=data,
+            link_index=L,
+            output_vel_repr=VelRepr.Inertial,
+        )
+
+        # Compute the Jacobian derivative of the parent link in inertial representation.
+        W_J̇_WL_W = js.link.jacobian_derivative(
+            model=model,
+            data=data,
+            link_index=L,
+            output_vel_repr=VelRepr.Inertial,
+        )
+
+    # =====================================================
+    # 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
+    # =====================================================
+
+    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:
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            O_X_W = F_X_W = Adjoint.from_transform(transform=W_H_F, inverse=True)
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                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
+
+        case VelRepr.Mixed:
+            W_H_F = transform(model=model, data=data, frame_index=frame_index)
+            W_H_FW = W_H_F.at[0:3, 0:3].set(jnp.eye(3))
+            FW_H_W = Transform.inverse(W_H_FW)
+            O_X_W = FW_X_W = Adjoint.from_transform(transform=FW_H_W)
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                FW_J_WF_FW = jacobian(
+                    model=model,
+                    data=data,
+                    frame_index=frame_index,
+                    output_vel_repr=VelRepr.Mixed,
+                )
+                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
+
+        case _:
+            raise ValueError(output_vel_repr)
+
+    O_J̇_WF_I = jnp.zeros(shape=(6, 6 + model.dofs()))
+    O_J̇_WF_I += O_Ẋ_W @ W_J_WL_W @ T
+    O_J̇_WF_I += O_X_W @ W_J̇_WL_W @ T
+    O_J̇_WF_I += O_X_W @ W_J_WL_W @ Ṫ
+
+    return O_J̇_WF_I

{jaxsim-0.4.1.dev26 → jaxsim-0.4.2.dev12}/src/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

{jaxsim-0.4.1.dev26 → jaxsim-0.4.2.dev12}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.1.dev26
+Version: 0.4.2.dev12
 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.1.dev26 → jaxsim-0.4.2.dev12}/tests/test_api_frame.py

@@ -5,6 +5,7 @@ import pytest
 
 import jaxsim.api as js
 from jaxsim import VelRepr
+from jaxsim.math.quaternion import Quaternion
 
 from . import utils_idyntree
 
@@ -168,3 +169,127 @@ def test_frame_jacobians(
         J_WL_js = js.frame.jacobian(model=model, data=data, frame_index=frame_index)
         J_WL_idt = kin_dyn.jacobian_frame(frame_name=frame_name)
         assert J_WL_js == pytest.approx(J_WL_idt, abs=1e-9)
+
+
+def test_frame_jacobian_derivative(
+    jaxsim_models_types: js.model.JaxSimModel,
+    velocity_representation: VelRepr,
+    prng_key: jax.Array,
+):
+    model = jaxsim_models_types
+
+    _, subkey = jax.random.split(prng_key, num=2)
+    data = js.data.random_model_data(
+        model=model, key=subkey, velocity_representation=velocity_representation
+    )
+
+    kin_dyn = utils_idyntree.build_kindyncomputations_from_jaxsim_model(
+        model=model, data=data
+    )
+
+    # Get all names of frames in the iDynTree model.
+    frame_names = [
+        frame.name
+        for frame in model.description.frames
+        if frame.name in kin_dyn.frame_names()
+    ]
+
+    # Skip some entry of models with many frames.
+    frame_names = [
+        name
+        for name in frame_names
+        if "skin" not in name or "laser" not in name or "depth" not in name
+    ]
+
+    frame_idxs = js.frame.names_to_idxs(model=model, frame_names=tuple(frame_names))
+
+    # ===============
+    # Test against AD
+    # ===============
+
+    # Get the generalized velocity.
+    I_ν = data.generalized_velocity()
+
+    # Compute J̇.
+    O_J̇_WF_I = jax.vmap(
+        lambda frame_index: js.frame.jacobian_derivative(
+            model=model, data=data, frame_index=frame_index
+        )
+    )(frame_idxs)
+
+    assert O_J̇_WF_I.shape == (len(frame_names), 6, 6 + model.dofs())
+
+    # Compute the plain Jacobian.
+    # This function will be used to compute the Jacobian derivative with AD.
+    def J(q, frame_idxs) -> jax.Array:
+        data_ad = js.data.JaxSimModelData.zero(
+            model=model, velocity_representation=data.velocity_representation
+        )
+
+        data_ad = data_ad.reset_base_position(base_position=q[:3])
+        data_ad = data_ad.reset_base_quaternion(base_quaternion=q[3:7])
+        data_ad = data_ad.reset_joint_positions(positions=q[7:])
+
+        O_J_ad_WF_I = jax.vmap(
+            lambda model, data, frame_index: js.frame.jacobian(
+                model=model, data=data, frame_index=frame_index
+            ),
+            in_axes=(None, None, 0),
+        )(model, data_ad, frame_idxs)
+
+        return O_J_ad_WF_I
+
+    def compute_q(data: js.data.JaxSimModelData) -> jax.Array:
+        q = jnp.hstack(
+            [
+                data.base_position(),
+                data.base_orientation(),
+                data.joint_positions(),
+            ]
+        )
+
+        return q
+
+    def compute_q̇(data: js.data.JaxSimModelData) -> jax.Array:
+        with data.switch_velocity_representation(VelRepr.Body):
+            B_ω_WB = data.base_velocity()[3:6]
+
+        with data.switch_velocity_representation(VelRepr.Mixed):
+            W_ṗ_B = data.base_velocity()[0:3]
+
+        W_Q̇_B = Quaternion.derivative(
+            quaternion=data.base_orientation(),
+            omega=B_ω_WB,
+            omega_in_body_fixed=True,
+            K=0.0,
+        ).squeeze()
+
+        q̇ = jnp.hstack([W_ṗ_B, W_Q̇_B, data.joint_velocities()])
+
+        return q̇
+
+    # Compute q and q̇.
+    q = compute_q(data)
+    q̇ = compute_q̇(data)
+
+    # Compute dJ/dt with AD.
+    dJ_dq = jax.jacfwd(J, argnums=0)(q, frame_idxs)
+    O_J̇_ad_WF_I = jnp.einsum("ijkq,q->ijk", dJ_dq, q̇)
+
+    assert O_J̇_WF_I == pytest.approx(expected=O_J̇_ad_WF_I)
+
+    # =====================
+    # Test against iDynTree
+    # =====================
+
+    # Compute the product J̇ν.
+    O_a_bias_WF = jax.vmap(
+        lambda O_J̇_WF_I, I_ν: O_J̇_WF_I @ I_ν,
+        in_axes=(0, None),
+    )(O_J̇_WF_I, I_ν)
+
+    # Compare the two computations.
+    for index, name in enumerate(frame_names):
+        J̇ν_idt = kin_dyn.frame_bias_acc(frame_name=name)
+        J̇ν_js = O_a_bias_WF[index]
+        assert J̇ν_js == pytest.approx(J̇ν_idt, abs=1e-9)