jaxsim 0.4.3.dev12__py3-none-any.whl → 0.4.3.dev18__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.3.dev12'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev12')
+__version__ = version = '0.4.3.dev18'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev18')

jaxsim/api/common.py

@@ -3,7 +3,7 @@ import contextlib
 import dataclasses
 import enum
 import functools
-from typing import ContextManager
+from collections.abc import Iterator
 
 import jax
 import jax.numpy as jnp
@@ -44,7 +44,7 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
     @contextlib.contextmanager
     def switch_velocity_representation(
         self, velocity_representation: VelRepr
-    ) -> ContextManager[Self]:
+    ) -> Iterator[Self]:
         """
         Context manager to temporarily switch the velocity representation.
 

jaxsim/api/contact.py

@@ -114,19 +114,24 @@ def collidable_point_forces(
 
 @jax.jit
 def collidable_point_dynamics(
-    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
-) -> tuple[jtp.Matrix, jtp.Matrix]:
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    link_forces: jtp.MatrixLike | None = None,
+) -> tuple[jtp.Matrix, dict[str, jtp.Array]]:
     r"""
-    Compute the 6D force applied to each collidable point and the corresponding
-    material deformation rate.
+    Compute the 6D force 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.
 
     Returns:
-        The 6D force applied to each collidable point and the corresponding
-        material deformation rate.
+        The 6D force applied to each collidable point and additional data based on the contact model configured:
+        - Soft: the material deformation rate.
+        - Rigid: nothing.
 
     Note:
         The material deformation rate is always returned in the mixed frame
@@ -138,7 +143,8 @@ def collidable_point_dynamics(
     # all collidable points belonging to the robot.
     W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
 
-    # Import privately the soft contacts classes.
+    # Import privately the contacts classes.
+    from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
     from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 
     # Build the soft contact model.
@@ -161,9 +167,31 @@ def collidable_point_dynamics(
             W_f_Ci, (CW_ṁ,) = jax.vmap(soft_contacts.compute_contact_forces)(
                 W_p_Ci, W_ṗ_Ci, data.state.contact.tangential_deformation
             )
+            aux_data = dict(m_dot=CW_ṁ)
+
+        case RigidContacts():
+            assert isinstance(model.contact_model, RigidContacts)
+            assert isinstance(data.state.contact, RigidContactsState)
+
+            # Build the contact model.
+            rigid_contacts = RigidContacts(
+                parameters=data.contacts_params, terrain=model.terrain
+            )
+
+            # Compute the 6D force expressed in the inertial frame and applied to each
+            # collidable point.
+            W_f_Ci, _ = rigid_contacts.compute_contact_forces(
+                position=W_p_Ci,
+                velocity=W_ṗ_Ci,
+                model=model,
+                data=data,
+                link_forces=link_forces,
+            )
+
+            aux_data = dict()
 
         case _:
-            raise ValueError("Invalid contact model {}".format(model.contact_model))
+            raise ValueError(f"Invalid contact model {model.contact_model}")
 
     # Convert the 6D forces to the active representation.
     f_Ci = jax.vmap(
@@ -175,7 +203,7 @@ def collidable_point_dynamics(
         )
     )(W_f_Ci)
 
-    return f_Ci, CW_ṁ
+    return f_Ci, aux_data
 
 
 @functools.partial(jax.jit, static_argnames=["link_names"])

jaxsim/api/data.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import functools
-from typing import Sequence
+from collections.abc import Sequence
 
 import jax
 import jax.numpy as jnp

jaxsim/api/frame.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Sequence
+from collections.abc import Sequence
 
 import jax
 import jax.numpy as jnp

jaxsim/api/joint.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Sequence
+from collections.abc import Sequence
 
 import jax
 import jax.numpy as jnp

jaxsim/api/link.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Sequence
+from collections.abc import Sequence
 
 import jax
 import jax.numpy as jnp

jaxsim/api/model.py

@@ -4,7 +4,8 @@ import copy
 import dataclasses
 import functools
 import pathlib
-from typing import Any, Sequence
+from collections.abc import Sequence
+from typing import Any
 
 import jax
 import jax.numpy as jnp
@@ -13,6 +14,7 @@ import rod
 from jax_dataclasses import Static
 
 import jaxsim.api as js
+import jaxsim.exceptions
 import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim.math import Adjoint, Cross
@@ -1889,6 +1891,8 @@ def step(
         and the new state of the integrator.
     """
 
+    from jaxsim.rbda.contacts.rigid import RigidContacts
+
     # Extract the integrator kwargs.
     # The following logic allows using integrators having kwargs colliding with the
     # kwargs of this step function.
@@ -1900,12 +1904,12 @@ def step(
 
     # Extract the initial resources.
     t0_ns = data.time_ns
-    state_x0 = data.state
+    state_t0 = data.state
     integrator_state_x0 = integrator_state
 
     # Step the dynamics forward.
-    state_xf, integrator_state_xf = integrator.step(
-        x0=state_x0,
+    state_tf, integrator_state_tf = integrator.step(
+        x0=state_t0,
         t0=jnp.array(t0_ns / 1e9).astype(float),
         dt=dt,
         params=integrator_state_x0,
@@ -1927,11 +1931,61 @@ def step(
         ),
     )
 
-    return (
+    data_tf = (
         # Store the new state of the model and the new time.
         data.replace(
-            state=state_xf,
+            state=state_tf,
             time_ns=t0_ns + jnp.array(dt * 1e9).astype(jnp.uint64),
-        ),
-        integrator_state_xf,
+        )
+    )
+
+    # Post process the simulation state, if needed.
+    match model.contact_model:
+
+        # Rigid contact models use an impact model that produces a discontinuous model velocity.
+        # Hence here we need to reset the velocity after each impact to guarantee that
+        # the linear velocity of the active collidable points is zero.
+        case RigidContacts():
+            # Raise runtime error for not supported case in which Rigid contacts and Baumgarte stabilization
+            # enabled are used with ForwardEuler integrator.
+            jaxsim.exceptions.raise_runtime_error_if(
+                condition=jnp.logical_and(
+                    isinstance(
+                        integrator,
+                        jaxsim.integrators.fixed_step.ForwardEuler
+                        | jaxsim.integrators.fixed_step.ForwardEulerSO3,
+                    ),
+                    jnp.array(
+                        [data_tf.contacts_params.K, data_tf.contacts_params.D]
+                    ).any(),
+                ),
+                msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
+            )
+
+            with data_tf.switch_velocity_representation(VelRepr.Mixed):
+                W_p_C = js.contact.collidable_point_positions(model, data_tf)
+                M = js.model.free_floating_mass_matrix(model, data_tf)
+                J_WC = js.contact.jacobian(model, data_tf)
+                px, py, _ = W_p_C.T
+                terrain_height = jax.vmap(model.terrain.height)(px, py)
+                inactive_collidable_points, _ = RigidContacts.detect_contacts(
+                    W_p_C=W_p_C,
+                    terrain_height=terrain_height,
+                )
+                BW_nu_post_impact = RigidContacts.compute_impact_velocity(
+                    data=data_tf,
+                    inactive_collidable_points=inactive_collidable_points,
+                    M=M,
+                    J_WC=J_WC,
+                )
+                data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
+                data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
+            # Restore the input velocity representation.
+            data_tf = data_tf.replace(
+                velocity_representation=data.velocity_representation, validate=False
+            )
+
+    return (
+        data_tf,
+        integrator_state_tf,
     )

jaxsim/api/ode.py

@@ -50,7 +50,7 @@ def wrap_system_dynamics_for_integration(
     # The wrapped dynamics will hold a reference of this object.
     model_closed = model.copy()
     data_closed = data.copy().replace(
-        state=js.ode_data.ODEState.zero(model=model_closed)
+        state=js.ode_data.ODEState.zero(model=model_closed, data=data)
     )
 
     def f(x: ODEState, t: Time, **kwargs_f) -> tuple[ODEState, dict[str, Any]]:
@@ -88,7 +88,7 @@ def system_velocity_dynamics(
     *,
     joint_forces: jtp.Vector | None = None,
     link_forces: jtp.Vector | None = None,
-) -> tuple[jtp.Vector, jtp.Vector, jtp.Matrix, dict[str, Any]]:
+) -> tuple[jtp.Vector, jtp.Vector, dict[str, Any]]:
     """
     Compute the dynamics of the system velocity.
 
@@ -102,18 +102,10 @@ def system_velocity_dynamics(
 
     Returns:
         A tuple containing the derivative of the base 6D velocity in inertial-fixed
-        representation, the derivative of the joint velocities, the derivative of
-        the material deformation, and the dictionary of auxiliary data returned by
-        the system dynamics evaluation.
+        representation, the derivative of the joint velocities, and auxiliary data
+        returned by the system dynamics evaluation.
     """
 
-    # Build joint torques if not provided.
-    τ = (
-        jnp.atleast_1d(joint_forces.squeeze())
-        if joint_forces is not None
-        else jnp.zeros_like(data.joint_positions())
-    ).astype(float)
-
     # Build link forces if not provided.
     # These forces are expressed in the frame corresponding to the velocity
     # representation of data.
@@ -123,6 +115,15 @@ def system_velocity_dynamics(
         else jnp.zeros((model.number_of_links(), 6))
     ).astype(float)
 
+    # We expect that the 6D forces included in the `link_forces` argument are expressed
+    # in the frame corresponding to the velocity representation of `data`.
+    references = js.references.JaxSimModelReferences.build(
+        model=model,
+        link_forces=O_f_L,
+        data=data,
+        velocity_representation=data.velocity_representation,
+    )
+
     # ======================
     # Compute contact forces
     # ======================
@@ -131,19 +132,17 @@ def system_velocity_dynamics(
     # with the terrain.
     W_f_Li_terrain = jnp.zeros_like(O_f_L).astype(float)
 
-    # Import privately the soft contacts classes.
-    from jaxsim.rbda.contacts.soft import SoftContactsState
-
-    # Initialize the derivative of the tangential deformation ṁ ∈ ℝ^{n_c × 3}.
-    assert isinstance(data.state.contact, SoftContactsState)
-    ṁ = jnp.zeros_like(data.state.contact.tangential_deformation).astype(float)
-
+    aux_data = {}
     if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
 
         # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
-        #  and the corresponding material deformation rates.
+        # along with contact-specific auxiliary states.
         with data.switch_velocity_representation(VelRepr.Inertial):
-            W_f_Ci, ṁ = js.contact.collidable_point_dynamics(model=model, data=data)
+            W_f_Ci, aux_data = js.contact.collidable_point_dynamics(
+                model=model,
+                data=data,
+                link_forces=references.link_forces(model=model, data=data),
+            )
 
         # 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
@@ -161,6 +160,74 @@ def system_velocity_dynamics(
 
         W_f_Li_terrain = mask.T @ W_f_Ci
 
+    # ===========================
+    # Compute system acceleration
+    # ===========================
+
+    # Compute the total link forces
+    with (
+        data.switch_velocity_representation(VelRepr.Inertial),
+        references.switch_velocity_representation(VelRepr.Inertial),
+    ):
+        references = references.apply_link_forces(
+            model=model,
+            data=data,
+            forces=W_f_Li_terrain,
+            additive=True,
+        )
+    # Get the link forces in the data representation
+    with references.switch_velocity_representation(data.velocity_representation):
+        f_L_total = references.link_forces(model=model, data=data)
+
+    # The following method always returns the inertial-fixed acceleration, and expects
+    # the link_forces expressed in the inertial frame.
+    W_v̇_WB, s̈ = system_acceleration(
+        model=model, data=data, joint_forces=joint_forces, link_forces=f_L_total
+    )
+
+    return W_v̇_WB, s̈, aux_data
+
+
+def system_acceleration(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    joint_forces: jtp.VectorLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
+) -> tuple[jtp.Vector, jtp.Vector]:
+    """
+    Compute the system acceleration in inertial-fixed representation.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        joint_forces: The joint forces to apply.
+        link_forces:
+            The 6D forces to apply to the links expressed in the same representation of data.
+
+    Returns:
+        A tuple containing the base 6D acceleration in inertial-fixed representation
+        and the joint accelerations.
+    """
+
+    # ====================
+    # Validate input data
+    # ====================
+
+    # Build link forces if not provided.
+    f_L = (
+        jnp.atleast_2d(link_forces.squeeze())
+        if link_forces is not None
+        else jnp.zeros((model.number_of_links(), 6))
+    ).astype(float)
+
+    # Build joint torques if not provided.
+    τ = (
+        jnp.atleast_1d(joint_forces.squeeze())
+        if joint_forces is not None
+        else jnp.zeros_like(data.joint_positions())
+    ).astype(float)
+
     # ====================
     # Enforce joint limits
     # ====================
@@ -198,29 +265,25 @@ def system_velocity_dynamics(
 
     references = js.references.JaxSimModelReferences.build(
         model=model,
-        joint_force_references=τ_total,
-        link_forces=O_f_L,
         data=data,
         velocity_representation=data.velocity_representation,
+        joint_force_references=τ_total,
+        link_forces=f_L,
     )
 
-    with references.switch_velocity_representation(VelRepr.Inertial):
-        W_f_L = references.link_forces(model=model, data=data)
-
-    # Compute the total external 6D forces applied to the links.
-    W_f_L_total = W_f_L + W_f_Li_terrain
-
     # - Joint accelerations: s̈ ∈ ℝⁿ
     # - Base inertial-fixed acceleration: W_v̇_WB = (W_p̈_B, W_ω̇_B) ∈ ℝ⁶
-    with data.switch_velocity_representation(velocity_representation=VelRepr.Inertial):
+    with (
+        data.switch_velocity_representation(velocity_representation=VelRepr.Inertial),
+        references.switch_velocity_representation(VelRepr.Inertial),
+    ):
         W_v̇_WB, s̈ = js.model.forward_dynamics_aba(
             model=model,
             data=data,
-            joint_forces=τ_total,
-            link_forces=W_f_L_total,
+            joint_forces=references.joint_force_references(),
+            link_forces=references.link_forces(),
         )
-
-    return W_v̇_WB, s̈, ṁ, dict()
+    return W_v̇_WB, s̈
 
 
 @jax.jit
@@ -291,14 +354,29 @@ def system_dynamics(
         by the system dynamics evaluation.
     """
 
+    from jaxsim.rbda.contacts.rigid import RigidContacts
+    from jaxsim.rbda.contacts.soft import SoftContacts
+
     # Compute the accelerations and the material deformation rate.
-    W_v̇_WB, s̈, ṁ, aux_dict = system_velocity_dynamics(
+    W_v̇_WB, s̈, aux_dict = system_velocity_dynamics(
         model=model,
         data=data,
         joint_forces=joint_forces,
         link_forces=link_forces,
     )
 
+    ode_state_kwargs = {}
+
+    match model.contact_model:
+        case SoftContacts():
+            ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
+
+        case RigidContacts():
+            pass
+
+        case _:
+            raise ValueError("Unable to determine contact state class prefix.")
+
     # Extract the velocities.
     W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
         model=model,
@@ -317,7 +395,7 @@ def system_dynamics(
         base_linear_velocity=W_v̇_WB[0:3],
         base_angular_velocity=W_v̇_WB[3:6],
         joint_velocities=s̈,
-        tangential_deformation=ṁ,
+        **ode_state_kwargs,
     )
 
     return ode_state_derivative, aux_dict

jaxsim/api/ode_data.py

@@ -6,6 +6,7 @@ import jax_dataclasses
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim.rbda import ContactsState
+from jaxsim.rbda.contacts.rigid import RigidContacts, RigidContactsState
 from jaxsim.rbda.contacts.soft import SoftContacts, SoftContactsState
 from jaxsim.utils import JaxsimDataclass
 
@@ -133,7 +134,7 @@ class ODEState(JaxsimDataclass):
         base_quaternion: jtp.Vector | None = None,
         base_linear_velocity: jtp.Vector | None = None,
         base_angular_velocity: jtp.Vector | None = None,
-        tangential_deformation: jtp.Matrix | None = None,
+        **kwargs,
     ) -> ODEState:
         """
         Build an `ODEState` from a `JaxSimModel`.
@@ -148,9 +149,7 @@ class ODEState(JaxsimDataclass):
                 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.
-            tangential_deformation:
-                The matrix of 3D tangential material deformations corresponding to
-                each collidable point.
+            kwargs: Additional arguments needed to build the contact state.
 
         Returns:
             The `ODEState` built from the `JaxSimModel`.
@@ -163,6 +162,7 @@ class ODEState(JaxsimDataclass):
         # Get the contact model from the `JaxSimModel`.
         match model.contact_model:
             case SoftContacts():
+                tangential_deformation = kwargs.get("tangential_deformation", None)
                 contact = SoftContactsState.build_from_jaxsim_model(
                     model=model,
                     **(
@@ -171,6 +171,8 @@ class ODEState(JaxsimDataclass):
                         else dict()
                     ),
                 )
+            case RigidContacts():
+                contact = RigidContactsState.build()
             case _:
                 raise ValueError("Unable to determine contact state class prefix.")
 
@@ -214,7 +216,7 @@ class ODEState(JaxsimDataclass):
 
         # Get the contact model from the `JaxSimModel`.
         match contact:
-            case SoftContactsState():
+            case SoftContactsState() | RigidContactsState():
                 pass
             case None:
                 contact = SoftContactsState.zero(model=model)
@@ -224,7 +226,7 @@ class ODEState(JaxsimDataclass):
         return ODEState(physics_model=physics_model_state, contact=contact)
 
     @staticmethod
-    def zero(model: js.model.JaxSimModel) -> ODEState:
+    def zero(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> ODEState:
         """
         Build a zero `ODEState` from a `JaxSimModel`.
 
@@ -235,7 +237,9 @@ class ODEState(JaxsimDataclass):
             A zero `ODEState` instance.
         """
 
-        model_state = ODEState.build(model=model)
+        model_state = ODEState.build(
+            model=model, contact=data.state.contact.zero(model=model)
+        )
 
         return model_state