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

jaxsim/mujoco/model.py

@@ -148,7 +148,7 @@ class MujocoModelHelper:
     def gravity(self) -> npt.NDArray:
         """Return the 3D gravity vector."""
 
-        return self.model.opt.gravity
+        return np.array([0, 0, self.model.gravity])
 
     # =========================
     # Methods for the base link

jaxsim/mujoco/utils.py

@@ -59,11 +59,11 @@ def mujoco_data_from_jaxsim(
     if jaxsim_model.floating_base():
 
         # Set the model position.
-        model_helper.set_base_position(position=np.array(jaxsim_data.base_position()))
+        model_helper.set_base_position(position=np.array(jaxsim_data.base_position))
 
         # Set the model orientation.
         model_helper.set_base_orientation(
-            orientation=np.array(jaxsim_data.base_orientation())
+            orientation=np.array(jaxsim_data.base_orientation)
         )
 
     # Set the joint positions.

jaxsim/parsers/kinematic_graph.py

@@ -952,9 +952,7 @@ class KinematicGraphTransforms:
         import jaxsim.math
 
         return np.array(
-            jaxsim.math.supported_joint_motion(joint_type, joint_position, joint_axis)[
-                0
-            ]
+            jaxsim.math.supported_joint_motion(joint_type, joint_position, joint_axis)
         )
 
     def find_parent_link_of_frame(self, name: str) -> str:

jaxsim/rbda/aba.py

@@ -4,7 +4,7 @@ import jaxlie
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint, Cross, StandardGravity
+from jaxsim.math import STANDARD_GRAVITY, Adjoint, Cross
 
 from . import utils
 
@@ -20,7 +20,7 @@ def aba(
     joint_velocities: jtp.VectorLike,
     joint_forces: jtp.VectorLike | None = None,
     link_forces: jtp.MatrixLike | None = None,
-    standard_gravity: jtp.FloatLike = StandardGravity,
+    standard_gravity: jtp.FloatLike = STANDARD_GRAVITY,
 ) -> tuple[jtp.Vector, jtp.Vector]:
     """
     Compute forward dynamics using the Articulated Body Algorithm (ABA).
@@ -85,13 +85,16 @@ def aba(
     W_X_B = W_H_B.adjoint()
     B_X_W = W_H_B.inverse().adjoint()
 
-    # Compute the parent-to-child adjoints and the motion subspaces of the joints.
+    # Compute the parent-to-child adjoints 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(
+    i_X_λi = model.kin_dyn_parameters.joint_transforms(
         joint_positions=s, base_transform=W_H_B.as_matrix()
     )
 
+    # Extract the joint motion subspaces.
+    S = model.kin_dyn_parameters.motion_subspaces
+
     # Allocate buffers.
     v = jnp.zeros(shape=(model.number_of_links(), 6, 1))
     c = jnp.zeros(shape=(model.number_of_links(), 6, 1))

jaxsim/rbda/collidable_points.py

@@ -1,23 +1,16 @@
 import jax
 import jax.numpy as jnp
-import jaxlie
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint, Skew
-
-from . import utils
+from jaxsim.math import Skew
 
 
 def collidable_points_pos_vel(
     model: js.model.JaxSimModel,
     *,
-    base_position: jtp.Vector,
-    base_quaternion: jtp.Vector,
-    joint_positions: jtp.Vector,
-    base_linear_velocity: jtp.Vector,
-    base_angular_velocity: jtp.Vector,
-    joint_velocities: jtp.Vector,
+    link_transforms: jtp.Matrix,
+    link_velocities: jtp.Matrix,
 ) -> tuple[jtp.Matrix, jtp.Matrix]:
     """
 
@@ -25,14 +18,8 @@ def collidable_points_pos_vel(
 
     Args:
         model: The model to consider.
-        base_position: The position of the base link.
-        base_quaternion: The quaternion of the base link.
-        joint_positions: The positions of the joints.
-        base_linear_velocity:
-            The linear velocity of the base link in inertial-fixed representation.
-        base_angular_velocity:
-            The angular velocity of the base link in inertial-fixed representation.
-        joint_velocities: The velocities of the joints.
+        link_transforms: The transforms from the world frame to each link.
+        link_velocities: The linear and angular velocities of each link.
 
     Returns:
         A tuple containing the position and linear velocity of the enabled collidable points.
@@ -54,95 +41,17 @@ def collidable_points_pos_vel(
     if len(indices_of_enabled_collidable_points) == 0:
         return jnp.array(0).astype(float), jnp.empty(0).astype(float)
 
-    W_p_B, W_Q_B, s, W_v_WB, ṡ, _, _, _, _, _ = utils.process_inputs(
-        model=model,
-        base_position=base_position,
-        base_quaternion=base_quaternion,
-        joint_positions=joint_positions,
-        base_linear_velocity=base_linear_velocity,
-        base_angular_velocity=base_angular_velocity,
-        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 base transform.
-    W_H_B = jaxlie.SE3.from_rotation_and_translation(
-        rotation=jaxlie.SO3(wxyz=W_Q_B),
-        translation=W_p_B,
-    )
-
-    # 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=W_H_B.as_matrix()
-    )
-
-    # Allocate buffer of transforms world -> link and initialize the base pose.
-    W_X_i = jnp.zeros(shape=(model.number_of_links(), 6, 6))
-    W_X_i = W_X_i.at[0].set(Adjoint.inverse(i_X_λi[0]))
-
-    # Allocate buffer of 6D inertial-fixed velocities and initialize the base velocity.
-    W_v_Wi = jnp.zeros(shape=(model.number_of_links(), 6))
-    W_v_Wi = W_v_Wi.at[0].set(W_v_WB)
-
-    # ====================
-    # Propagate kinematics
-    # ====================
-
-    PropagateTransformsCarry = tuple[jtp.Matrix, jtp.Matrix]
-    propagate_transforms_carry: PropagateTransformsCarry = (W_X_i, W_v_Wi)
-
-    def propagate_kinematics(
-        carry: PropagateTransformsCarry, i: jtp.Int
-    ) -> tuple[PropagateTransformsCarry, None]:
-
-        ii = i - 1
-        W_X_i, W_v_Wi = carry
-
-        # Compute the parent to child 6D transform.
-        λi_X_i = Adjoint.inverse(adjoint=i_X_λi[i])
-
-        # Compute the world to child 6D transform.
-        W_Xi_i = W_X_i[λ[i]] @ λi_X_i
-        W_X_i = W_X_i.at[i].set(W_Xi_i)
-
-        # Propagate the 6D velocity.
-        W_vi_Wi = W_v_Wi[λ[i]] + W_X_i[i] @ (S[i] * ṡ[ii]).squeeze()
-        W_v_Wi = W_v_Wi.at[i].set(W_vi_Wi)
-
-        return (W_X_i, W_v_Wi), None
-
-    (W_X_i, W_v_Wi), _ = (
-        jax.lax.scan(
-            f=propagate_kinematics,
-            init=propagate_transforms_carry,
-            xs=jnp.arange(start=1, stop=model.number_of_links()),
-        )
-        if model.number_of_links() > 1
-        else [(W_X_i, W_v_Wi), None]
-    )
-
-    # ==================================================
-    # Compute position and velocity of collidable points
-    # ==================================================
-
     def process_point_kinematics(
         Li_p_C: jtp.Vector, parent_body: jtp.Int
     ) -> tuple[jtp.Vector, jtp.Vector]:
 
         # Compute the position of the collidable point.
-        W_p_Ci = (
-            Adjoint.to_transform(adjoint=W_X_i[parent_body]) @ jnp.hstack([Li_p_C, 1])
-        )[0:3]
+        W_p_Ci = (link_transforms[parent_body] @ jnp.hstack([Li_p_C, 1]))[0:3]
 
         # Compute the linear part of the mixed velocity Ci[W]_v_{W,Ci}.
         CW_vl_WCi = (
             jnp.block([jnp.eye(3), -Skew.wedge(vector=W_p_Ci).squeeze()])
-            @ W_v_Wi[parent_body].squeeze()
+            @ link_velocities[parent_body].squeeze()
         )
 
         return W_p_Ci, CW_vl_WCi

jaxsim/rbda/contacts/__init__.py

@@ -1,13 +1,5 @@
-from . import relaxed_rigid, rigid, soft, visco_elastic
+from . import relaxed_rigid
 from .common import ContactModel, ContactsParams
 from .relaxed_rigid import RelaxedRigidContacts, RelaxedRigidContactsParams
-from .rigid import RigidContacts, RigidContactsParams
-from .soft import SoftContacts, SoftContactsParams
-from .visco_elastic import ViscoElasticContacts, ViscoElasticContactsParams
 
-ContactParamsTypes = (
-    SoftContactsParams
-    | RigidContactsParams
-    | RelaxedRigidContactsParams
-    | ViscoElasticContactsParams
-)
+ContactParamsTypes = RelaxedRigidContactsParams

jaxsim/rbda/contacts/common.py

@@ -9,7 +9,6 @@ import jax.numpy as jnp
 import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
-from jaxsim.api.common import ModelDataWithVelocityRepresentation
 from jaxsim.utils import JaxsimDataclass
 
 try:
@@ -135,143 +134,6 @@ class ContactModel(JaxsimDataclass):
 
         pass
 
-    def compute_link_contact_forces(
-        self,
-        model: js.model.JaxSimModel,
-        data: js.data.JaxSimModelData,
-        **kwargs,
-    ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
-        """
-        Compute the link contact forces.
-
-        Args:
-            model: The robot model considered by the contact model.
-            data: The data of the considered model.
-            **kwargs: Optional additional arguments, specific to the contact model.
-
-        Returns:
-            A tuple containing as first element the 6D contact force applied to the
-            links and expressed in the frame of the velocity representation of data,
-            and as second element a dictionary of optional additional information.
-        """
-
-        # Compute the contact forces expressed in the inertial frame.
-        # This function, contrarily to `compute_contact_forces`, already handles how
-        # the optional kwargs should be passed to the specific contact models.
-        W_f_C, aux_dict = js.contact.collidable_point_dynamics(
-            model=model, data=data, **kwargs
-        )
-
-        # Compute the 6D forces applied to the links equivalent to the forces applied
-        # to the frames associated to the collidable points.
-        with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
-
-            W_f_L = self.link_forces_from_contact_forces(
-                model=model, data=data, contact_forces=W_f_C
-            )
-
-        # Store the link forces in the references object for easy conversion.
-        references = js.references.JaxSimModelReferences.build(
-            model=model,
-            data=data,
-            link_forces=W_f_L,
-            velocity_representation=jaxsim.VelRepr.Inertial,
-        )
-
-        # Convert the link forces to the frame corresponding to the velocity
-        # representation of data.
-        with references.switch_velocity_representation(data.velocity_representation):
-            f_L = references.link_forces(model=model, data=data)
-
-        return f_L, aux_dict
-
-    @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.
-        f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
-
-        # Get the pose of the enabled collidable points.
-        W_H_C = js.contact.transforms(model=model, data=data)[
-            indices_of_enabled_collidable_points
-        ]
-
-        # Convert the contact forces to inertial-fixed representation.
-        W_f_C = jax.vmap(
-            lambda f_C, W_H_C: (
-                ModelDataWithVelocityRepresentation.other_representation_to_inertial(
-                    array=f_C,
-                    other_representation=data.velocity_representation,
-                    transform=W_H_C,
-                    is_force=True,
-                )
-            )
-        )(f_C, W_H_C)
-
-        # 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
-
-        # Compute the link transforms.
-        W_H_L = (
-            js.model.forward_kinematics(model=model, data=data)
-            if data.velocity_representation is not jaxsim.VelRepr.Inertial
-            else jnp.zeros(shape=(model.number_of_links(), 4, 4))
-        )
-
-        # Convert the inertial-fixed link forces to the velocity representation of data.
-        f_L = jax.vmap(
-            lambda W_f_L, W_H_L: (
-                ModelDataWithVelocityRepresentation.inertial_to_other_representation(
-                    array=W_f_L,
-                    other_representation=data.velocity_representation,
-                    transform=W_H_L,
-                    is_force=True,
-                )
-            )
-        )(W_f_L, W_H_L)
-
-        return f_L
-
     @classmethod
     def zero_state_variables(cls, model: js.model.JaxSimModel) -> dict[str, jtp.Array]:
         """

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+import functools
 from collections.abc import Callable
 from typing import Any
 
@@ -13,6 +14,7 @@ import jaxsim.api as js
 import jaxsim.rbda.contacts
 import jaxsim.typing as jtp
 from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
+from jaxsim.terrain.terrain import Terrain
 
 from . import common
 
@@ -263,7 +265,7 @@ class RelaxedRigidContacts(common.ContactModel):
                 Optional `(n_joints,)` vector of joint forces.
 
         Returns:
-            A tuple containing as first element the computed contact forces.
+            A tuple containing as first element the computed contact forces in inertial representation.
         """
 
         link_forces = jnp.atleast_2d(
@@ -306,20 +308,17 @@ class RelaxedRigidContacts(common.ContactModel):
         W_H_C = js.contact.transforms(model=model, data=data)
 
         with (
-            references.switch_velocity_representation(VelRepr.Mixed),
             data.switch_velocity_representation(VelRepr.Mixed),
+            references.switch_velocity_representation(VelRepr.Mixed),
         ):
-
-            BW_ν = data.generalized_velocity()
+            BW_ν = data.generalized_velocity
 
             BW_ν̇_free = jnp.hstack(
                 js.ode.system_acceleration(
                     model=model,
                     data=data,
                     link_forces=references.link_forces(model=model, data=data),
-                    joint_force_references=references.joint_force_references(
-                        model=model
-                    ),
+                    joint_torques=references.joint_force_references(model=model),
                 )
             )
 
@@ -342,7 +341,7 @@ class RelaxedRigidContacts(common.ContactModel):
             model=model,
             position_constraint=position_constraint,
             velocity_constraint=velocity,
-            parameters=data.contacts_params,
+            parameters=model.contacts_params,
         )
 
         # Compute the Delassus matrix and the free mixed linear acceleration of
@@ -426,7 +425,7 @@ class RelaxedRigidContacts(common.ContactModel):
 
         # Initialize the optimized forces with a linear Hunt/Crossley model.
         init_params = jax.vmap(
-            lambda p, v: jaxsim.rbda.contacts.SoftContacts.hunt_crossley_contact_model(
+            lambda p, v: self._hunt_crossley_contact_model(
                 position=p,
                 velocity=v,
                 terrain=model.terrain,
@@ -603,3 +602,149 @@ class RelaxedRigidContacts(common.ContactModel):
         )
 
         return a_ref, jnp.diag(R), K, D
+
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=("terrain",))
+    def _hunt_crossley_contact_model(
+        position: jtp.VectorLike,
+        velocity: jtp.VectorLike,
+        tangential_deformation: jtp.VectorLike,
+        terrain: Terrain,
+        K: jtp.FloatLike,
+        D: jtp.FloatLike,
+        mu: jtp.FloatLike,
+        p: jtp.FloatLike = 0.5,
+        q: jtp.FloatLike = 0.5,
+    ) -> tuple[jtp.Vector, jtp.Vector]:
+        """
+        Compute the contact force using the Hunt/Crossley model.
+
+        Args:
+            position: The position of the collidable point.
+            velocity: The velocity of the collidable point.
+            tangential_deformation: The material deformation of the collidable point.
+            terrain: The terrain model.
+            K: The stiffness parameter.
+            D: The damping parameter of the soft contacts model.
+            mu: The static friction coefficient.
+            p:
+                The exponent p corresponding to the damping-related non-linearity
+                of the Hunt/Crossley model.
+            q:
+                The exponent q corresponding to the spring-related non-linearity
+                of the Hunt/Crossley model
+
+        Returns:
+            A tuple containing the computed contact force and the derivative of the
+            material deformation.
+        """
+
+        # Convert the input vectors to arrays.
+        W_p_C = jnp.array(position, dtype=float).squeeze()
+        W_ṗ_C = jnp.array(velocity, dtype=float).squeeze()
+        m = jnp.array(tangential_deformation, dtype=float).squeeze()
+
+        # Use symbol for the static friction.
+        μ = mu
+
+        # Compute the penetration depth, its rate, and the considered terrain normal.
+        δ, δ̇, n̂ = common.compute_penetration_data(p=W_p_C, v=W_ṗ_C, terrain=terrain)
+
+        # There are few operations like computing the norm of a vector with zero length
+        # or computing the square root of zero that are problematic in an AD context.
+        # To avoid these issues, we introduce a small tolerance ε to their arguments
+        # and make sure that we do not check them against zero directly.
+        ε = jnp.finfo(float).eps
+
+        # Compute the powers of the penetration depth.
+        # Inject ε to address AD issues in differentiating the square root when
+        #  p and q are fractional.
+        δp = jnp.power(δ + ε, p)
+        δq = jnp.power(δ + ε, q)
+
+        # ========================
+        # Compute the normal force
+        # ========================
+
+        # Non-linear spring-damper model (Hunt/Crossley model).
+        # This is the force magnitude along the direction normal to the terrain.
+        force_normal_mag = (K * δp) * δ + (D * δq) * δ̇
+
+        # Depending on the magnitude of δ̇, the normal force could be negative.
+        force_normal_mag = jnp.maximum(0.0, force_normal_mag)
+
+        # Compute the 3D linear force in C[W] frame.
+        f_normal = force_normal_mag * n̂
+
+        # ============================
+        # Compute the tangential force
+        # ============================
+
+        # Extract the tangential component of the velocity.
+        v_tangential = W_ṗ_C - jnp.dot(W_ṗ_C, n̂) * n̂
+
+        # Extract the normal and tangential components of the material deformation.
+        m_normal = jnp.dot(m, n̂) * n̂
+        m_tangential = m - jnp.dot(m, n̂) * n̂
+
+        # Compute the tangential force in the sticking case.
+        # Using the tangential component of the material deformation should not be
+        # necessary if the sticking-slipping transition occurs in a terrain area
+        # with a locally constant normal. However, this assumption is not true in
+        # general, especially for highly uneven terrains.
+        f_tangential = -((K * δp) * m_tangential + (D * δq) * v_tangential)
+
+        # Detect the contact type (sticking or slipping).
+        # Note that if there is no contact, sticking is set to True, and this detail
+        # is exploited in the computation of the `contact_status` variable.
+        sticking = jnp.logical_or(
+            δ <= 0, f_tangential.dot(f_tangential) <= (μ * force_normal_mag) ** 2
+        )
+
+        # Compute the direction of the tangential force.
+        # To prevent dividing by zero, we use a switch statement.
+        norm = jaxsim.math.safe_norm(f_tangential)
+        f_tangential_direction = f_tangential / (
+            norm + jnp.finfo(float).eps * (norm == 0)
+        )
+
+        # Project the tangential force to the friction cone if slipping.
+        f_tangential = jnp.where(
+            sticking,
+            f_tangential,
+            jnp.minimum(μ * force_normal_mag, norm) * f_tangential_direction,
+        )
+
+        # Set the tangential force to zero if there is no contact.
+        f_tangential = jnp.where(δ <= 0, jnp.zeros(3), f_tangential)
+
+        # =====================================
+        # Compute the material deformation rate
+        # =====================================
+
+        # Compute the derivative of the material deformation.
+        # Note that we included an additional relaxation of `m_normal` in the
+        # sticking case, so that the normal deformation that could have accumulated
+        # from a previous slipping phase can relax to zero.
+        ṁ_no_contact = -(K / D) * m
+        ṁ_sticking = v_tangential - (K / D) * m_normal
+        ṁ_slipping = -(f_tangential + (K * δp) * m_tangential) / (D * δq)
+
+        # Compute the contact status:
+        # 0: slipping
+        # 1: sticking
+        # 2: no contact
+        contact_status = sticking.astype(int)
+        contact_status += (δ <= 0).astype(int)
+
+        # Select the right material deformation rate depending on the contact status.
+        ṁ = jax.lax.select_n(contact_status, ṁ_slipping, ṁ_sticking, ṁ_no_contact)
+
+        # ==========================================
+        # Compute and return the final contact force
+        # ==========================================
+
+        # Sum the normal and tangential forces.
+        CW_fl = f_normal + f_tangential
+
+        return CW_fl, ṁ

jaxsim/rbda/crba.py

@@ -30,13 +30,16 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
     # 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.
+    # Compute the parent-to-child adjoints 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(
+    i_X_λi = model.kin_dyn_parameters.joint_transforms(
         joint_positions=s, base_transform=jnp.eye(4)
     )
 
+    # Extract the joint motion subspaces.
+    S = model.kin_dyn_parameters.motion_subspaces
+
     # Allocate the buffer of transforms link -> base.
     i_X_0 = jnp.zeros(shape=(model.number_of_links(), 6, 6))
     i_X_0 = i_X_0.at[0].set(jnp.eye(6))