jaxsim 0.4.2.dev19__tar.gz → 0.4.2.dev28__tar.gz

{jaxsim-0.4.2.dev19 → jaxsim-0.4.2.dev28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.2.dev19
+Version: 0.4.2.dev28
 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 → jaxsim-0.4.2.dev28}/src/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.dev28'
+__version_tuple__ = version_tuple = (0, 4, 2, 'dev28')

{jaxsim-0.4.2.dev19 → jaxsim-0.4.2.dev28}/src/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-0.4.2.dev19 → jaxsim-0.4.2.dev28}/src/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-0.4.2.dev19 → jaxsim-0.4.2.dev28}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.2.dev19
+Version: 0.4.2.dev28
 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.dev28/tests/test_api_contact.py

@@ -0,0 +1,147 @@
+import jax
+import jax.numpy as jnp
+import pytest
+import rod
+
+import jaxsim.api as js
+from jaxsim import VelRepr
+
+
+def test_contact_kinematics(
+    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,
+    )
+
+    # =====
+    # Tests
+    # =====
+
+    # Compute the pose of the implicit contact frame associated to the collidable points
+    # and the transforms of all links.
+    W_H_C = js.contact.transforms(model=model, data=data)
+    W_H_L = js.model.forward_kinematics(model=model, data=data)
+
+    # Check that the orientation of the implicit contact frame matches with the
+    # orientation of the link to which the contact point is attached.
+    for contact_idx, index_of_parent_link in enumerate(
+        model.kin_dyn_parameters.contact_parameters.body
+    ):
+        assert W_H_C[contact_idx, 0:3, 0:3] == pytest.approx(
+            W_H_L[index_of_parent_link][0:3, 0:3]
+        )
+
+    # Check that the origin of the implicit contact frame is located over the
+    # collidable point.
+    W_p_C = js.contact.collidable_point_positions(model=model, data=data)
+    assert W_p_C == pytest.approx(W_H_C[:, 0:3, 3])
+
+    # Compute the velocity of the collidable point.
+    # This quantity always matches with the linear component of the mixed 6D velocity
+    # of the implicit frame associated to the collidable point.
+    W_ṗ_C = js.contact.collidable_point_velocities(model=model, data=data)
+
+    # Compute the velocity of the collidable point using the contact Jacobian.
+    ν = data.generalized_velocity()
+    CW_J_WC = js.contact.jacobian(model=model, data=data, output_vel_repr=VelRepr.Mixed)
+    CW_vl_WC = jnp.einsum("c6g,g->c6", CW_J_WC, ν)[:, 0:3]
+
+    # Compare the two velocities.
+    assert W_ṗ_C == pytest.approx(CW_vl_WC)
+
+
+def test_contact_jacobian_derivative(
+    jaxsim_models_floating_base: js.model.JaxSimModel,
+    velocity_representation: VelRepr,
+    prng_key: jax.Array,
+):
+
+    model = jaxsim_models_floating_base
+
+    _, subkey = jax.random.split(prng_key, num=2)
+    data = js.data.random_model_data(
+        model=model,
+        key=subkey,
+        velocity_representation=velocity_representation,
+    )
+
+    # =====
+    # Tests
+    # =====
+
+    # Extract the parent link names and the poses of the contact points.
+    parent_link_names = js.link.idxs_to_names(
+        model=model, link_indices=model.kin_dyn_parameters.contact_parameters.body
+    )
+    W_p_Ci = model.kin_dyn_parameters.contact_parameters.point
+
+    # Load the model in ROD.
+    rod_model = rod.Sdf.load(sdf=model.built_from, is_urdf=True).model
+
+    # Add dummy frames on the contact points.
+    for idx, (link_name, W_p_C) in enumerate(
+        zip(parent_link_names, W_p_Ci, strict=True)
+    ):
+        rod_model.add_frame(
+            frame=rod.Frame(
+                name=f"contact_point_{idx}",
+                attached_to=link_name,
+                pose=rod.Pose(
+                    relative_to=link_name, pose=jnp.zeros(shape=(6,)).at[0:3].set(W_p_C)
+                ),
+            ),
+        )
+
+    # Rebuild the JaxSim model.
+    model_with_frames = js.model.JaxSimModel.build_from_model_description(
+        model_description=rod_model
+    )
+    model_with_frames = js.model.reduce(
+        model=model_with_frames, considered_joints=model.joint_names()
+    )
+
+    # Rebuild the JaxSim data.
+    data_with_frames = js.data.JaxSimModelData.build(
+        model=model_with_frames,
+        base_position=data.base_position(),
+        base_quaternion=data.base_orientation(dcm=False),
+        joint_positions=data.joint_positions(),
+        base_linear_velocity=data.base_velocity()[0:3],
+        base_angular_velocity=data.base_velocity()[3:6],
+        joint_velocities=data.joint_velocities(),
+        velocity_representation=velocity_representation,
+    )
+
+    # Extract the indexes of the frames attached to the contact points.
+    frame_idxs = js.frame.names_to_idxs(
+        model=model_with_frames,
+        frame_names=(
+            f"contact_point_{idx}" for idx in list(range(len(parent_link_names)))
+        ),
+    )
+
+    # Check that the number of frames is correct.
+    assert len(frame_idxs) == len(parent_link_names)
+
+    # Compute the contact Jacobian derivative.
+    J̇_WC = js.contact.jacobian_derivative(
+        model=model_with_frames, data=data_with_frames
+    )
+
+    # Compute the contact Jacobian derivative using frames.
+    J̇_WF = jax.vmap(
+        js.frame.jacobian_derivative,
+        in_axes=(None, None),
+    )(model_with_frames, data, frame_index=frame_idxs)
+
+    # Compare the two Jacobians.
+    assert J̇_WC == pytest.approx(J̇_WF)

jaxsim-0.4.2.dev19/tests/test_api_contact.py

@@ -1,58 +0,0 @@
-import jax
-import jax.numpy as jnp
-import pytest
-
-import jaxsim.api as js
-from jaxsim import VelRepr
-
-
-def test_contact_kinematics(
-    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,
-    )
-
-    # =====
-    # Tests
-    # =====
-
-    # Compute the pose of the implicit contact frame associated to the collidable points
-    # and the transforms of all links.
-    W_H_C = js.contact.transforms(model=model, data=data)
-    W_H_L = js.model.forward_kinematics(model=model, data=data)
-
-    # Check that the orientation of the implicit contact frame matches with the
-    # orientation of the link to which the contact point is attached.
-    for contact_idx, index_of_parent_link in enumerate(
-        model.kin_dyn_parameters.contact_parameters.body
-    ):
-        assert W_H_C[contact_idx, 0:3, 0:3] == pytest.approx(
-            W_H_L[index_of_parent_link][0:3, 0:3]
-        )
-
-    # Check that the origin of the implicit contact frame is located over the
-    # collidable point.
-    W_p_C = js.contact.collidable_point_positions(model=model, data=data)
-    assert W_p_C == pytest.approx(W_H_C[:, 0:3, 3])
-
-    # Compute the velocity of the collidable point.
-    # This quantity always matches with the linear component of the mixed 6D velocity
-    # of the implicit frame associated to the collidable point.
-    W_ṗ_C = js.contact.collidable_point_velocities(model=model, data=data)
-
-    # Compute the velocity of the collidable point using the contact Jacobian.
-    ν = data.generalized_velocity()
-    CW_J_WC = js.contact.jacobian(model=model, data=data, output_vel_repr=VelRepr.Mixed)
-    CW_vl_WC = jnp.einsum("c6g,g->c6", CW_J_WC, ν)[:, 0:3]
-
-    # Compare the two velocities.
-    assert W_ṗ_C == pytest.approx(CW_vl_WC)