jaxsim 0.6.2.dev2__py3-none-any.whl → 0.6.2.dev102__py3-none-any.whl

jaxsim/__init__.py

@@ -114,5 +114,5 @@ del _get_default_logging_level
 del _is_editable
 
 from . import terrain  # isort:skip
-from . import api, integrators, logging, math, rbda
+from . import api, logging, math, rbda
 from .api.common import VelRepr

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.2.dev2'
-__version_tuple__ = version_tuple = (0, 6, 2, 'dev2')
+__version__ = version = '0.6.2.dev102'
+__version_tuple__ = version_tuple = (0, 6, 2, 'dev102')

jaxsim/api/__init__.py

@@ -1,13 +1,15 @@
 from . import common  # isort:skip
 from . import model, data  # isort:skip
 from . import (
+    actuation_model,
     com,
     contact,
+    contact_model,
     frame,
+    integrators,
     joint,
     kin_dyn_parameters,
     link,
     ode,
-    ode_data,
     references,
 )

jaxsim/api/actuation_model.py

@@ -0,0 +1,96 @@
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+
+
+def compute_resultant_torques(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    joint_force_references: jtp.Vector | None = None,
+) -> jtp.Vector:
+    """
+    Compute the resultant torques acting on the joints.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        joint_force_references: The joint force references to apply.
+
+    Returns:
+        The resultant torques acting on the joints.
+    """
+
+    # Build joint torques if not provided.
+    τ_references = (
+        jnp.atleast_1d(joint_force_references.squeeze())
+        if joint_force_references is not None
+        else jnp.zeros_like(data.joint_positions)
+    ).astype(float)
+
+    # ====================
+    # Enforce joint limits
+    # ====================
+
+    τ_position_limit = jnp.zeros_like(τ_references).astype(float)
+
+    if model.dofs() > 0:
+
+        # Stiffness and damper parameters for the joint position limits.
+        k_j = jnp.array(
+            model.kin_dyn_parameters.joint_parameters.position_limit_spring
+        ).astype(float)
+        d_j = jnp.array(
+            model.kin_dyn_parameters.joint_parameters.position_limit_damper
+        ).astype(float)
+
+        # Compute the joint position limit violations.
+        lower_violation = jnp.clip(
+            data.joint_positions
+            - model.kin_dyn_parameters.joint_parameters.position_limits_min,
+            max=0.0,
+        )
+
+        upper_violation = jnp.clip(
+            data.joint_positions
+            - model.kin_dyn_parameters.joint_parameters.position_limits_max,
+            min=0.0,
+        )
+
+        # Compute the joint position limit torque.
+        τ_position_limit -= jnp.diag(k_j) @ (lower_violation + upper_violation)
+
+        τ_position_limit -= (
+            jnp.positive(τ_position_limit) * jnp.diag(d_j) @ data.joint_velocities
+        )
+
+    # ====================
+    # Joint friction model
+    # ====================
+
+    τ_friction = jnp.zeros_like(τ_references).astype(float)
+
+    if model.dofs() > 0:
+
+        # Static and viscous joint friction parameters
+        kc = jnp.array(
+            model.kin_dyn_parameters.joint_parameters.friction_static
+        ).astype(float)
+        kv = jnp.array(
+            model.kin_dyn_parameters.joint_parameters.friction_viscous
+        ).astype(float)
+
+        # Compute the joint friction torque.
+        τ_friction = -(
+            jnp.diag(kc) @ jnp.sign(data.joint_velocities)
+            + jnp.diag(kv) @ data.joint_velocities
+        )
+
+    # ===============================
+    # Compute the total joint forces.
+    # ===============================
+
+    τ_total = τ_references + τ_friction + τ_position_limit
+
+    return τ_total

jaxsim/api/com.py

@@ -26,8 +26,8 @@ def com_position(
 
     m = js.model.total_mass(model=model)
 
-    W_H_L = js.model.forward_kinematics(model=model, data=data)
-    W_H_B = data.base_transform()
+    W_H_L = data._link_transforms
+    W_H_B = data._base_transform
     B_H_W = jaxsim.math.Transform.inverse(transform=W_H_B)
 
     def B_p̃_LCoM(i) -> jtp.Vector:
@@ -98,7 +98,7 @@ def centroidal_momentum(
         and :math:`C = B` if the active velocity representation is body-fixed.
     """
 
-    ν = data.generalized_velocity()
+    ν = data.generalized_velocity
     G_J = centroidal_momentum_jacobian(model=model, data=data)
 
     return G_J @ ν
@@ -134,7 +134,7 @@ def centroidal_momentum_jacobian(
         model=model, data=data, output_vel_repr=VelRepr.Body
     )
 
-    W_H_B = data.base_transform()
+    W_H_B = data._base_transform
     B_H_W = jaxsim.math.Transform.inverse(W_H_B)
 
     W_p_CoM = com_position(model=model, data=data)
@@ -172,7 +172,7 @@ def locked_centroidal_spatial_inertia(
     with data.switch_velocity_representation(VelRepr.Body):
         B_Mbb_B = js.model.locked_spatial_inertia(model=model, data=data)
 
-    W_H_B = data.base_transform()
+    W_H_B = data._base_transform
     W_p_CoM = com_position(model=model, data=data)
 
     match data.velocity_representation:
@@ -213,7 +213,7 @@ def average_centroidal_velocity(
         and :math:`[C] = [B]` if the active velocity representation is body-fixed.
     """
 
-    ν = data.generalized_velocity()
+    ν = data.generalized_velocity
     G_J = average_centroidal_velocity_jacobian(model=model, data=data)
 
     return G_J @ ν
@@ -269,7 +269,7 @@ def bias_acceleration(
     """
 
     # Compute the pose of all links with forward kinematics.
-    W_H_L = js.model.forward_kinematics(model=model, data=data)
+    W_H_L = data._link_transforms
 
     # Compute the bias acceleration of all links by zeroing the generalized velocity
     # in the active representation.
@@ -411,7 +411,7 @@ def bias_acceleration(
         case VelRepr.Body:
 
             GB_Xf_W = jaxsim.math.Adjoint.from_transform(
-                transform=data.base_transform().at[0:3].set(W_p_CoM)
+                transform=data._base_transform.at[0:3].set(W_p_CoM)
             ).T
 
             GB_ḣ_bias = GB_Xf_W @ W_ḣ_bias

jaxsim/api/contact.py

@@ -42,12 +42,8 @@ def collidable_point_kinematics(
 
         W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
             model=model,
-            base_position=data.base_position(),
-            base_quaternion=data.base_orientation(dcm=False),
-            joint_positions=data.joint_positions(model=model),
-            base_linear_velocity=data.base_velocity()[0:3],
-            base_angular_velocity=data.base_velocity()[3:6],
-            joint_velocities=data.joint_velocities(model=model),
+            link_transforms=data._link_transforms,
+            link_velocities=data._link_velocities,
         )
 
     return W_p_Ci, W_ṗ_Ci
@@ -95,143 +91,6 @@ def collidable_point_velocities(
     return W_ṗ_Ci
 
 
-@jax.jit
-@js.common.named_scope
-def collidable_point_forces(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
-    link_forces: jtp.MatrixLike | None = None,
-    joint_force_references: jtp.VectorLike | None = None,
-    **kwargs,
-) -> jtp.Matrix:
-    """
-    Compute the 6D forces applied to each collidable point.
-
-    Args:
-        model: The model to consider.
-        data: The data of the considered model.
-        link_forces:
-            The 6D external forces to apply to the links expressed in the same
-            representation of data.
-        joint_force_references:
-            The joint force references to apply to the joints.
-        kwargs: Additional keyword arguments to pass to the active contact model.
-
-    Returns:
-        The 6D forces applied to each collidable point expressed in the frame
-        corresponding to the active representation.
-    """
-
-    f_Ci, _ = collidable_point_dynamics(
-        model=model,
-        data=data,
-        link_forces=link_forces,
-        joint_force_references=joint_force_references,
-        **kwargs,
-    )
-
-    return f_Ci
-
-
-@jax.jit
-@js.common.named_scope
-def collidable_point_dynamics(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
-    link_forces: jtp.MatrixLike | None = None,
-    joint_force_references: jtp.VectorLike | None = None,
-    **kwargs,
-) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
-    r"""
-    Compute the 6D force applied to each enabled collidable point.
-
-    Args:
-        model: The model to consider.
-        data: The data of the considered model.
-        link_forces:
-            The 6D external forces to apply to the links expressed in the same
-            representation of data.
-        joint_force_references:
-            The joint force references to apply to the joints.
-        kwargs: Additional keyword arguments to pass to the active contact model.
-
-    Returns:
-        The 6D force applied to each enabled collidable point and additional data based
-        on the contact model configured:
-        - Soft: the material deformation rate.
-        - Rigid: no additional data.
-        - QuasiRigid: no additional data.
-
-    Note:
-        The material deformation rate is always returned in the mixed frame
-        `C[W] = ({}^W \mathbf{p}_C, [W])`. This is convenient for integration purpose.
-        Instead, the 6D forces are returned in the active representation.
-    """
-
-    # Build the common kw arguments to pass to the computation of the contact forces.
-    common_kwargs = dict(
-        link_forces=link_forces,
-        joint_force_references=joint_force_references,
-    )
-
-    # Build the additional kwargs to pass to the computation of the contact forces.
-    match model.contact_model:
-
-        case contacts.SoftContacts():
-
-            kwargs_contact_model = {}
-
-        case contacts.RigidContacts():
-
-            kwargs_contact_model = common_kwargs | kwargs
-
-        case contacts.RelaxedRigidContacts():
-
-            kwargs_contact_model = common_kwargs | kwargs
-
-        case contacts.ViscoElasticContacts():
-
-            kwargs_contact_model = common_kwargs | dict(dt=model.time_step) | kwargs
-
-        case _:
-            raise ValueError(f"Invalid contact model: {model.contact_model}")
-
-    # Compute the contact forces with the active contact model.
-    W_f_C, aux_data = model.contact_model.compute_contact_forces(
-        model=model,
-        data=data,
-        **kwargs_contact_model,
-    )
-
-    # Compute the transforms of the implicit frames `C[L] = (W_p_C, [L])`
-    # associated to the enabled collidable point.
-    # In inertial-fixed representation, the computation of these transforms
-    # is not necessary and the conversion below becomes a no-op.
-
-    # Get the indices of the enabled collidable points.
-    indices_of_enabled_collidable_points = (
-        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
-    )
-
-    W_H_C = (
-        js.contact.transforms(model=model, data=data)
-        if data.velocity_representation is not VelRepr.Inertial
-        else jnp.stack([jnp.eye(4)] * len(indices_of_enabled_collidable_points))
-    )
-
-    # Convert the 6D forces to the active representation.
-    f_Ci = jax.vmap(
-        lambda W_f_C, W_H_C: data.inertial_to_other_representation(
-            array=W_f_C,
-            other_representation=data.velocity_representation,
-            transform=W_H_C,
-            is_force=True,
-        )
-    )(W_f_C, W_H_C)
-
-    return f_Ci, aux_data
-
-
 @functools.partial(jax.jit, static_argnames=["link_names"])
 @js.common.named_scope
 def in_contact(
@@ -305,11 +164,7 @@ def estimate_good_soft_contacts_parameters(
 def estimate_good_contact_parameters(
     model: js.model.JaxSimModel,
     *,
-    standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
     static_friction_coefficient: jtp.FloatLike = 0.5,
-    number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
-    damping_ratio: jtp.FloatLike = 1.0,
-    max_penetration: jtp.FloatLike | None = None,
     **kwargs,
 ) -> jaxsim.rbda.contacts.ContactParamsTypes:
     """
@@ -317,15 +172,7 @@ def estimate_good_contact_parameters(
 
     Args:
         model: The model to consider.
-        standard_gravity: The standard gravity constant.
         static_friction_coefficient: The static friction coefficient.
-        number_of_active_collidable_points_steady_state:
-            The number of active collidable points in steady state supporting
-            the weight of the robot.
-        damping_ratio: The damping ratio.
-        max_penetration:
-            The maximum penetration allowed in steady state when the robot is
-            supported by the configured number of active collidable points.
         kwargs:
             Additional model-specific parameters passed to the builder method of
             the parameters class.
@@ -343,82 +190,8 @@ def estimate_good_contact_parameters(
         specific application.
     """
 
-    def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
-        """
-        Displacement between the CoM and the lowest collidable point using zero
-        joint positions.
-        """
-
-        zero_data = js.data.JaxSimModelData.build(
-            model=model,
-            contacts_params=jaxsim.rbda.contacts.SoftContactsParams(),
-        )
-
-        W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
-
-        if model.floating_base():
-            W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
-            return 2 * (W_pz_CoM - W_pz_C.min())
-
-        return 2 * W_pz_CoM
-
-    max_δ = (
-        max_penetration
-        if max_penetration is not None
-        # Consider as default a 0.5% of the model height.
-        else 0.005 * estimate_model_height(model=model)
-    )
-
-    nc = number_of_active_collidable_points_steady_state
-
     match model.contact_model:
 
-        case contacts.SoftContacts():
-            assert isinstance(model.contact_model, contacts.SoftContacts)
-
-            parameters = contacts.SoftContactsParams.build_default_from_jaxsim_model(
-                model=model,
-                standard_gravity=standard_gravity,
-                static_friction_coefficient=static_friction_coefficient,
-                max_penetration=max_δ,
-                number_of_active_collidable_points_steady_state=nc,
-                damping_ratio=damping_ratio,
-                **kwargs,
-            )
-
-        case contacts.ViscoElasticContacts():
-            assert isinstance(model.contact_model, contacts.ViscoElasticContacts)
-
-            parameters = (
-                contacts.ViscoElasticContactsParams.build_default_from_jaxsim_model(
-                    model=model,
-                    standard_gravity=standard_gravity,
-                    static_friction_coefficient=static_friction_coefficient,
-                    max_penetration=max_δ,
-                    number_of_active_collidable_points_steady_state=nc,
-                    damping_ratio=damping_ratio,
-                    **kwargs,
-                )
-            )
-
-        case contacts.RigidContacts():
-            assert isinstance(model.contact_model, contacts.RigidContacts)
-
-            # Disable Baumgarte stabilization by default since it does not play
-            # well with the forward Euler integrator.
-            K = kwargs.get("K", 0.0)
-
-            parameters = contacts.RigidContactsParams.build(
-                mu=static_friction_coefficient,
-                **(
-                    dict(
-                        K=K,
-                        D=2 * jnp.sqrt(K),
-                    )
-                    | kwargs
-                ),
-            )
-
         case contacts.RelaxedRigidContacts():
             assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
 
@@ -463,9 +236,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
     )[indices_of_enabled_collidable_points]
 
     # Get the transforms of the parent link of all collidable points.
-    W_H_L = js.model.forward_kinematics(model=model, data=data)[
-        parent_link_idx_of_enabled_collidable_points
-    ]
+    W_H_L = data._link_transforms[parent_link_idx_of_enabled_collidable_points]
 
     L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
         indices_of_enabled_collidable_points
@@ -615,7 +386,10 @@ def jacobian_derivative(
     ]
 
     # Get the transforms of all the parent links.
-    W_H_Li = js.model.forward_kinematics(model=model, data=data)
+    W_H_Li = data._link_transforms
+
+    # Get the link velocities.
+    W_v_WLi = data._link_velocities
 
     # =====================================================
     # Compute quantities to adjust the input representation
@@ -643,9 +417,9 @@ def jacobian_derivative(
             Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_W)
 
         case VelRepr.Body:
-            W_H_B = data.base_transform()
+            W_H_B = data._base_transform
             W_X_B = Adjoint.from_transform(transform=W_H_B)
-            B_v_WB = data.base_velocity()
+            B_v_WB = data.base_velocity
             B_vx_WB = Cross.vx(B_v_WB)
             W_Ẋ_B = W_X_B @ B_vx_WB
 
@@ -653,10 +427,10 @@ def jacobian_derivative(
             Ṫ = compute_Ṫ(model=model, Ẋ=W_Ẋ_B)
 
         case VelRepr.Mixed:
-            W_H_B = data.base_transform()
+            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_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
@@ -676,27 +450,16 @@ def jacobian_derivative(
         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 enabled 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,
         parent_link_idx: jtp.Int,
-        CW_J_WC_BW: jtp.Matrix,
         W_H_L: jtp.Matrix,
     ) -> jtp.Matrix:
 
@@ -711,9 +474,7 @@ def jacobian_derivative(
                 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_v_WC = W_v_WLi[parent_link_idx]
                 W_vx_WC = Cross.vx(W_v_WC)
                 O_Ẋ_W = C_Ẋ_W = -C_X_W @ W_vx_WC  # noqa: F841
 
@@ -723,8 +484,7 @@ def jacobian_derivative(
                 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()
+                CW_v_WC = CW_X_W @ W_v_WLi[parent_link_idx]
                 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
@@ -739,8 +499,8 @@ def jacobian_derivative(
 
         return O_J̇_WC_I
 
-    O_J̇_WC = jax.vmap(compute_O_J̇_WC_I, in_axes=(0, 0, 0, None))(
-        L_p_Ci, parent_link_idx_of_enabled_collidable_points, CW_J_WC_BW, W_H_Li
+    O_J̇_WC = jax.vmap(compute_O_J̇_WC_I, in_axes=(0, 0, None))(
+        L_p_Ci, parent_link_idx_of_enabled_collidable_points, W_H_Li
     )
 
     return O_J̇_WC

jaxsim/api/contact_model.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+
+
+@jax.jit
+@js.common.named_scope
+def link_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_torques: jtp.VectorLike | None = None,
+) -> jtp.Matrix:
+    """
+    Compute the 6D contact forces of all links of the model in inertial representation.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_forces:
+            The 6D external forces to apply to the links expressed in inertial representation
+        joint_torques:
+            The joint torques acting on the joints.
+
+    Returns:
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
+        expressed in inertial representation.
+    """
+
+    # Compute the contact forces for each collidable point with the active contact model.
+    W_f_C, _ = model.contact_model.compute_contact_forces(
+        model=model,
+        data=data,
+        link_forces=link_forces,
+        joint_force_references=joint_torques,
+    )
+
+    # Compute the 6D forces applied to the links equivalent to the forces applied
+    # to the frames associated to the collidable points.
+    W_f_L = link_forces_from_contact_forces(
+        model=model, data=data, contact_forces=W_f_C
+    )
+
+    return W_f_L
+
+
+@staticmethod
+def link_forces_from_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    contact_forces: jtp.MatrixLike,
+) -> jtp.Matrix:
+    """
+    Compute the link forces from the contact forces.
+
+    Args:
+        model: The robot model considered by the contact model.
+        data: The data of the considered model.
+        contact_forces: The contact forces computed by the contact model.
+
+    Returns:
+        The 6D contact forces applied to the links and expressed in the frame of
+        the velocity representation of data.
+    """
+
+    # Get the object storing the contact parameters of the model.
+    contact_parameters = model.kin_dyn_parameters.contact_parameters
+
+    # Extract the indices corresponding to the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    # Convert the contact forces to a JAX array.
+    W_f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
+
+    # Construct the vector defining the parent link index of each collidable point.
+    # We use this vector to sum the 6D forces of all collidable points rigidly
+    # attached to the same link.
+    parent_link_index_of_collidable_points = jnp.array(
+        contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    # Create the mask that associate each collidable point to their parent link.
+    # We use this mask to sum the collidable points to the right link.
+    mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+        model.number_of_links()
+    )
+
+    # Sum the forces of all collidable points rigidly attached to a body.
+    # Since the contact forces W_f_C are expressed in the world frame,
+    # we don't need any coordinate transformation.
+    W_f_L = mask.T @ W_f_C
+
+    return W_f_L