jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev2__py3-none-any.whl

jaxsim/api/ode.py

@@ -2,64 +2,99 @@ from typing import Any, Protocol
 
 import jax
 import jax.numpy as jnp
-import jaxlie
 
-import jaxsim.physics.algos.soft_contacts
+import jaxsim.api as js
+import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim import VelRepr, integrators
-from jaxsim.integrators.common import Time
-from jaxsim.math.quaternion import Quaternion
-from jaxsim.physics.algos.soft_contacts import SoftContactsState
-from jaxsim.physics.model.physics_model_state import PhysicsModelState
-from jaxsim.simulation.ode_data import ODEState
+from jaxsim.integrators import Time
+from jaxsim.math import Quaternion
+from jaxsim.rbda import contacts
 
-from . import contact as Contact
-from . import data as Data
-from . import model as Model
+from .common import VelRepr
+from .ode_data import ODEState
 
 
 class SystemDynamicsFromModelAndData(Protocol):
+    """
+    Protocol defining the signature of a function computing the system dynamics
+    given a model and data object.
+    """
+
     def __call__(
         self,
-        model: Model.JaxSimModel,
-        data: Data.JaxSimModelData,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
         **kwargs: dict[str, Any],
-    ) -> tuple[ODEState, dict[str, Any]]: ...
+    ) -> tuple[ODEState, dict[str, Any]]:
+        """
+        Compute the system dynamics given a model and data object.
+
+        Args:
+            model: The model to consider.
+            data: The data of the considered model.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            A tuple with an `ODEState` object storing in each of its attributes the
+            corresponding derivative, and the dictionary of auxiliary data returned
+            by the system dynamics evaluation.
+        """
+
+        pass
 
 
 def wrap_system_dynamics_for_integration(
-    model: Model.JaxSimModel,
-    data: Data.JaxSimModelData,
     *,
     system_dynamics: SystemDynamicsFromModelAndData,
-    **kwargs,
+    **kwargs: dict[str, Any],
 ) -> jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]:
     """
-    Wrap generic system dynamics operating on `JaxSimModel` and `JaxSimModelData`
-    for integration with `jaxsim.integrators`.
+    Wrap the system dynamics considered by JaxSim integrators in a generic
+    `f(x, t, **u, **parameters)` function.
 
     Args:
-        model: The model to consider.
-        data: The data of the considered model.
         system_dynamics: The system dynamics to wrap.
         **kwargs: Additional kwargs to close over the system dynamics.
 
     Returns:
-        The system dynamics closed over the model, the data, and the additional kwargs.
+        The system dynamics closed over the additional kwargs to be used by
+        JaxSim integrators.
     """
 
-    # We allow to close `system_dynamics` over additional kwargs.
-    kwargs_closed = kwargs
-
-    def f(x: ODEState, t: Time, **kwargs) -> tuple[ODEState, dict[str, Any]]:
-
-        # Close f over the `data` parameter.
-        with data.editable(validate=True) as data_rw:
+    # Close `system_dynamics` over additional kwargs.
+    # Similarly to what done in `jaxsim.api.model.step`, to be future-proof, we use the
+    # following logic to allow the caller to close over arguments having the same name
+    # of the ones used in the `wrap_system_dynamics_for_integration` function.
+    kwargs = kwargs.copy() if kwargs is not None else {}
+    colliding_system_dynamics_kwargs = kwargs.pop("system_dynamics_kwargs", {})
+    system_dynamics_kwargs = kwargs | colliding_system_dynamics_kwargs
+
+    # Remove `model` and `data` for backward compatibility.
+    # It's no longer necessary to close over them at this stage, as this is always
+    # done in `jaxsim.api.model.step`.
+    # We can remove the following lines in a few releases.
+    _ = system_dynamics_kwargs.pop("data", None)
+    _ = system_dynamics_kwargs.pop("model", None)
+
+    # Create the function with the signature expected by our generic integrators.
+    # Note that our system dynamics is time independent.
+    def f(x: ODEState, t: Time, **kwargs_f) -> tuple[ODEState, dict[str, Any]]:
+
+        # Get the data and model objects from the kwargs.
+        data_f = kwargs_f.pop("data")
+        model_f = kwargs_f.pop("model")
+
+        # Update the state and time stored inside data.
+        with data_f.editable(validate=True) as data_rw:
             data_rw.state = x
-            data_rw.time_ns = jnp.array(t * 1e9).astype(jnp.uint64)
 
-        # Close f over the `model` parameter.
-        return system_dynamics(model=model, data=data_rw, **kwargs_closed | kwargs)
+        # Evaluate the system dynamics, allowing to override the kwargs originally
+        # passed when the closure was created.
+        return system_dynamics(
+            model=model_f,
+            data=data_rw,
+            **(system_dynamics_kwargs | kwargs_f),
+        )
 
     f: jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]
     return f
@@ -71,105 +106,217 @@ def wrap_system_dynamics_for_integration(
 
 
 @jax.jit
+@js.common.named_scope
 def system_velocity_dynamics(
-    model: Model.JaxSimModel,
-    data: Data.JaxSimModelData,
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.Vector | None = None,
-    external_forces: jtp.Vector | None = None,
-) -> tuple[jtp.Vector, jtp.Vector, jtp.Matrix, dict[str, Any]]:
+    link_forces: jtp.Vector | None = None,
+    joint_force_references: jtp.Vector | None = None,
+) -> tuple[jtp.Vector, jtp.Vector, dict[str, Any]]:
     """
     Compute the dynamics of the system velocity.
 
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint forces to apply.
-        external_forces: The external forces to apply to the links.
+        link_forces:
+            The 6D forces to apply to the links expressed in the frame corresponding to
+            the velocity representation of `data`.
+        joint_force_references: The joint force references to apply.
 
     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 evalutation.
+        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 external forces if not provided
-    f_ext = (
-        jnp.atleast_2d(external_forces.squeeze())
-        if external_forces is not None
+    # Build link forces if not provided.
+    # These forces are expressed in the frame corresponding to the velocity
+    # representation of data.
+    O_f_L = (
+        jnp.atleast_2d(link_forces.squeeze())
+        if link_forces is not None
         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,
+        joint_force_references=joint_force_references,
+        data=data,
+        velocity_representation=data.velocity_representation,
+    )
+
     # ======================
     # Compute contact forces
     # ======================
 
     # Initialize the 6D forces W_f ∈ ℝ^{n_L × 6} applied to links due to contact
     # with the terrain.
-    W_f_Li_terrain = jnp.zeros_like(f_ext).astype(float)
-
-    # Initialize the 6D contact forces W_f ∈ ℝ^{n_c × 3} applied to collidable points,
-    # expressed in the world frame.
-    W_f_Ci = None
-
-    # Initialize the derivative of the tangential deformation ṁ ∈ ℝ^{n_c × 3}.
-    ṁ = jnp.zeros_like(data.state.soft_contacts.tangential_deformation).astype(float)
-
-    if len(model.physics_model.gc.body) > 0:
-        # Compute the position and linear velocities (mixed representation) of
-        # all collidable points belonging to the robot.
-        W_p_Ci, W_ṗ_Ci = Contact.collidable_point_kinematics(model=model, data=data)
-
-        # Compute the 3D forces applied to each collidable point.
-        W_f_Ci, ṁ = jax.vmap(
-            lambda p, ṗ, m: jaxsim.physics.algos.soft_contacts.SoftContacts(
-                parameters=data.soft_contacts_params, terrain=model.terrain
-            ).contact_model(position=p, velocity=ṗ, tangential_deformation=m)
-        )(W_p_Ci, W_ṗ_Ci, data.state.soft_contacts.tangential_deformation.T)
-
-        # Sum the forces of all collidable points rigidly attached to a body.
-        # Since the contact forces W_f_Ci are expressed in the world frame,
-        # we don't need any coordinate transformation.
-        W_f_Li_terrain = jax.vmap(
-            lambda nc: (
-                jnp.vstack(
-                    jnp.equal(
-                        np.array(model.physics_model.gc.body, dtype=int), nc
-                    ).astype(int)
-                )
-                * W_f_Ci
-            ).sum(axis=0)
-        )(jnp.arange(model.number_of_links()))
+    W_f_L_terrain = jnp.zeros_like(O_f_L).astype(float)
+
+    # Initialize a dictionary of auxiliary data.
+    # This dictionary is used to store additional data computed by the contact model.
+    aux_data = {}
+
+    if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
+
+        with (
+            data.switch_velocity_representation(VelRepr.Inertial),
+            references.switch_velocity_representation(VelRepr.Inertial),
+        ):
+
+            # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
+            # along with contact-specific auxiliary states.
+            W_f_C, aux_data = js.contact.collidable_point_dynamics(
+                model=model,
+                data=data,
+                link_forces=references.link_forces(model=model, data=data),
+                joint_force_references=references.joint_force_references(model=model),
+            )
+
+            # Compute the 6D forces applied to the links equivalent to the forces applied
+            # to the frames associated to the collidable points.
+            W_f_L_terrain = model.contact_model.link_forces_from_contact_forces(
+                model=model,
+                data=data,
+                contact_forces=W_f_C,
+            )
+
+    # ===========================
+    # Compute system acceleration
+    # ===========================
+
+    # Compute the total link forces.
+    with (
+        data.switch_velocity_representation(VelRepr.Inertial),
+        references.switch_velocity_representation(VelRepr.Inertial),
+    ):
+
+        # Sum the contact forces just computed with the link forces applied by the user.
+        references = references.apply_link_forces(
+            model=model,
+            data=data,
+            forces=W_f_L_terrain,
+            additive=True,
+        )
+
+        # Get the link forces in inertial-fixed representation.
+        f_L_total = references.link_forces(model=model, data=data)
+
+        # Compute the system acceleration in inertial-fixed representation.
+        # This representation is useful for integration purpose.
+        W_v̇_WB, s̈ = system_acceleration(
+            model=model,
+            data=data,
+            joint_force_references=joint_force_references,
+            link_forces=f_L_total,
+        )
+
+    return W_v̇_WB, s̈, aux_data
+
+
+def system_acceleration(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_force_references: jtp.VectorLike | None = None,
+) -> tuple[jtp.Vector, jtp.Vector]:
+    """
+    Compute the system acceleration in the active representation.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_forces:
+            The 6D forces to apply to the links expressed in the same
+            velocity representation of data.
+        joint_force_references: The joint force references to apply.
+
+    Returns:
+        A tuple containing the base 6D acceleration in the active 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.
+    τ_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
     # ====================
 
-    # TODO: enforce joint limits
-    τ_position_limit = jnp.zeros_like(τ).astype(float)
+    τ_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.state.physics_model.joint_positions
+            - model.kin_dyn_parameters.joint_parameters.position_limits_min,
+            max=0.0,
+        )
+
+        upper_violation = jnp.clip(
+            data.state.physics_model.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.state.physics_model.joint_velocities
+        )
 
     # ====================
     # Joint friction model
     # ====================
 
-    τ_friction = jnp.zeros_like(τ).astype(float)
+    τ_friction = jnp.zeros_like(τ_references).astype(float)
 
     if model.dofs() > 0:
-        # Static and viscous joint friction parameters
-        kc = jnp.array(list(model.physics_model._joint_friction_static.values()))
-        kv = jnp.array(list(model.physics_model._joint_friction_viscous.values()))
 
-        # Compute the joint friction torque
+        # 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.state.physics_model.joint_positions)
+            jnp.diag(kc) @ jnp.sign(data.state.physics_model.joint_velocities)
             + jnp.diag(kv) @ data.state.physics_model.joint_velocities
         )
 
@@ -177,28 +324,40 @@ def system_velocity_dynamics(
     # Compute forward dynamics
     # ========================
 
-    # Compute the total joint forces
-    τ_total = τ + τ_friction + τ_position_limit
+    # Compute the total joint forces.
+    τ_total = τ_references + τ_friction + τ_position_limit
 
-    # Compute the total external 6D forces applied to the links
-    W_f_L_total = f_ext + W_f_Li_terrain
+    # Store the link forces in a references object.
+    references = js.references.JaxSimModelReferences.build(
+        model=model,
+        data=data,
+        velocity_representation=data.velocity_representation,
+        link_forces=f_L,
+    )
 
+    # Compute forward dynamics.
+    #
     # - Joint accelerations: s̈ ∈ ℝⁿ
-    # - Base inertial-fixed acceleration: W_v̇_WB = (W_p̈_B, W_ω̇_B) ∈ ℝ⁶
-    with data.switch_velocity_representation(velocity_representation=VelRepr.Inertial):
-        W_v̇_WB, s̈ = Model.forward_dynamics_aba(
-            model=model,
-            data=data,
-            joint_forces=τ_total,
-            external_forces=W_f_L_total,
-        )
+    # - Base acceleration: v̇_WB ∈ ℝ⁶
+    #
+    # Note that ABA returns the base acceleration in the velocity representation
+    # stored in the `data` object.
+    v̇_WB, s̈ = js.model.forward_dynamics_aba(
+        model=model,
+        data=data,
+        joint_forces=τ_total,
+        link_forces=references.link_forces(model=model, data=data),
+    )
 
-    return W_v̇_WB, s̈, ṁ.T, dict()
+    return v̇_WB, s̈
 
 
 @jax.jit
+@js.common.named_scope
 def system_position_dynamics(
-    model: Model.JaxSimModel, data: Data.JaxSimModelData
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[jtp.Vector, jtp.Vector, jtp.Vector]:
     """
     Compute the dynamics of the system position.
@@ -206,14 +365,16 @@ def system_position_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
+        baumgarte_quaternion_regularization:
+            The Baumgarte regularization coefficient for adjusting the quaternion norm.
 
     Returns:
         A tuple containing the derivative of the base position, the derivative of the
         base quaternion, and the derivative of the joint positions.
     """
 
-    ṡ = data.state.physics_model.joint_velocities
-    W_Q_B = data.state.physics_model.base_quaternion
+    ṡ = data.joint_velocities(model=model)
+    W_Q_B = data.base_orientation(dcm=False)
 
     with data.switch_velocity_representation(velocity_representation=VelRepr.Mixed):
         W_ṗ_B = data.base_velocity()[0:3]
@@ -225,18 +386,21 @@ def system_position_dynamics(
         quaternion=W_Q_B,
         omega=W_ω_WB,
         omega_in_body_fixed=False,
+        K=baumgarte_quaternion_regularization,
     ).squeeze()
 
     return W_ṗ_B, W_Q̇_B, ṡ
 
 
 @jax.jit
+@js.common.named_scope
 def system_dynamics(
-    model: Model.JaxSimModel,
-    data: Data.JaxSimModelData,
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
     *,
-    joint_forces: jtp.Vector | None = None,
-    external_forces: jtp.Vector | None = None,
+    link_forces: jtp.Vector | None = None,
+    joint_force_references: jtp.Vector | None = None,
+    baumgarte_quaternion_regularization: jtp.FloatLike = 1.0,
 ) -> tuple[ODEState, dict[str, Any]]:
     """
     Compute the dynamics of the system.
@@ -244,8 +408,13 @@ def system_dynamics(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        joint_forces: The joint forces to apply.
-        external_forces: The external forces to apply to the links.
+        link_forces:
+            The 6D forces to apply to the links expressed in the frame corresponding to
+            the velocity representation of `data`.
+        joint_force_references: The joint force references to apply.
+        baumgarte_quaternion_regularization:
+            The Baumgarte regularization coefficient used to adjust the norm of the
+            quaternion (only used in integrators not operating on the SO(3) manifold).
 
     Returns:
         A tuple with an `ODEState` object storing in each of its attributes the
@@ -254,31 +423,53 @@ def system_dynamics(
     """
 
     # 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,
-        external_forces=external_forces,
+        joint_force_references=joint_force_references,
+        link_forces=link_forces,
     )
 
+    # Initialize the dictionary storing the derivative of the additional state variables
+    # that extend the state vector of the integrated ODE system.
+    extended_ode_state = {}
+
+    match model.contact_model:
+
+        case contacts.SoftContacts():
+            extended_ode_state["tangential_deformation"] = aux_dict["m_dot"]
+
+        case contacts.ViscoElasticContacts():
+
+            extended_ode_state["tangential_deformation"] = jnp.zeros_like(
+                data.state.extended["tangential_deformation"]
+            )
+
+        case contacts.RigidContacts() | contacts.RelaxedRigidContacts():
+            pass
+
+        case _:
+            raise ValueError(f"Invalid contact model: {model.contact_model}")
+
     # Extract the velocities.
-    W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(model=model, data=data)
+    W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
+        model=model,
+        data=data,
+        baumgarte_quaternion_regularization=baumgarte_quaternion_regularization,
+    )
 
     # Create an ODEState object populated with the derivative of each leaf.
     # Our integrators, operating on generic pytrees, will be able to handle it
     # automatically as state derivative.
-    ode_state_derivative = ODEState.build(
-        physics_model_state=PhysicsModelState.build(
-            base_position=W_ṗ_B,
-            base_quaternion=W_Q̇_B,
-            joint_positions=ṡ,
-            base_linear_velocity=W_v̇_WB[0:3],
-            base_angular_velocity=W_v̇_WB[3:6],
-            joint_velocities=s̈,
-        ),
-        soft_contacts_state=SoftContactsState.build(
-            tangential_deformation=ṁ,
-        ),
+    ode_state_derivative = ODEState.build_from_jaxsim_model(
+        model=model,
+        base_position=W_ṗ_B,
+        base_quaternion=W_Q̇_B,
+        joint_positions=ṡ,
+        base_linear_velocity=W_v̇_WB[0:3],
+        base_angular_velocity=W_v̇_WB[3:6],
+        joint_velocities=s̈,
+        **extended_ode_state,
     )
 
     return ode_state_derivative, aux_dict