jaxsim 0.4.3.dev229__tar.gz → 0.4.3.dev242__tar.gz

{jaxsim-0.4.3.dev229 → jaxsim-0.4.3.dev242}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev229
+Version: 0.4.3.dev242
 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.3.dev229 → jaxsim-0.4.3.dev242}/examples/jaxsim_as_multibody_dynamics_library.ipynb

@@ -728,7 +728,7 @@
     "Key methodologies in this area may involve the Delassus matrix:\n",
     "\n",
     "$$\n",
-    "\\Psi(\\mathbf{q}) = J_{W,C}^T(\\mathbf{q}) \\, M(\\mathbf{q})^{-1} \\, J_{W,C}^T(\\mathbf{q})\n",
+    "\\Psi(\\mathbf{q}) = J_{W,C}(\\mathbf{q}) \\, M(\\mathbf{q})^{-1} \\, J_{W,C}^T(\\mathbf{q})\n",
     "$$\n",
     "\n",
     "or the linear acceleration of a contact point:\n",
@@ -780,7 +780,7 @@
     "    J̇l_WC = js.contact.jacobian_derivative(model=model, data=data)[:, 0:3, :]\n",
     "\n",
     "# Compute the Delassus matrix.\n",
-    "Ψ = jnp.vstack(Jl_WC) @ jnp.linalg.lstsq(BW_M, jnp.vstack(J̇l_WC).T)[0]\n",
+    "Ψ = jnp.vstack(Jl_WC) @ jnp.linalg.lstsq(BW_M, jnp.vstack(Jl_WC).T)[0]\n",
     "print(f\"Ψ:     shape={Ψ.shape}\")\n",
     "\n",
     "# Compute the transforms of the mixed frames implicitly associated\n",

{jaxsim-0.4.3.dev229 → jaxsim-0.4.3.dev242}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev229'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev229')
+__version__ = version = '0.4.3.dev242'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev242')

{jaxsim-0.4.3.dev229 → jaxsim-0.4.3.dev242}/src/jaxsim/api/contact.py

@@ -98,6 +98,7 @@ def collidable_point_forces(
     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.
@@ -110,6 +111,7 @@ def collidable_point_forces(
             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
@@ -121,6 +123,7 @@ def collidable_point_forces(
         data=data,
         link_forces=link_forces,
         joint_force_references=joint_force_references,
+        **kwargs,
     )
 
     return f_Ci
@@ -132,7 +135,8 @@ def collidable_point_dynamics(
     data: js.data.JaxSimModelData,
     link_forces: jtp.MatrixLike | None = None,
     joint_force_references: jtp.VectorLike | None = None,
-) -> tuple[jtp.Matrix, dict[str, jtp.Array]]:
+    **kwargs,
+) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
     r"""
     Compute the 6D force applied to each collidable point.
 
@@ -144,6 +148,7 @@ def collidable_point_dynamics(
             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 collidable point and additional data based
@@ -158,86 +163,46 @@ def collidable_point_dynamics(
         Instead, the 6D forces are returned in the active representation.
     """
 
-    # Build the soft contact model.
+    # 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():
-            assert isinstance(model.contact_model, contacts.SoftContacts)
 
-            # Compute the 6D force expressed in the inertial frame and applied to each
-            # collidable point, and the corresponding material deformation rate.
-            # Note that the material deformation rate is always returned in the mixed frame
-            # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
-            W_f_Ci, (CW_ṁ,) = model.contact_model.compute_contact_forces(
-                model=model, data=data
-            )
-
-            # Create the dictionary of auxiliary data.
-            # This contact model considers the material deformation as additional state
-            # of the ODE system. We need to pass its dynamics to the integrator.
-            aux_data = dict(m_dot=CW_ṁ)
+            kwargs_contact_model = {}
 
         case contacts.RigidContacts():
-            assert isinstance(model.contact_model, contacts.RigidContacts)
 
-            # Compute the 6D force expressed in the inertial frame and applied to each
-            # collidable point.
-            W_f_Ci, _ = model.contact_model.compute_contact_forces(
-                model=model,
-                data=data,
-                link_forces=link_forces,
-                joint_force_references=joint_force_references,
-            )
-
-            aux_data = dict()
+            kwargs_contact_model = common_kwargs | kwargs
 
         case contacts.RelaxedRigidContacts():
-            assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
 
-            # Compute the 6D force expressed in the inertial frame and applied to each
-            # collidable point.
-            W_f_Ci, _ = model.contact_model.compute_contact_forces(
-                model=model,
-                data=data,
-                link_forces=link_forces,
-                joint_force_references=joint_force_references,
-            )
-
-            aux_data = dict()
+            kwargs_contact_model = common_kwargs | kwargs
 
         case contacts.ViscoElasticContacts():
-            assert isinstance(model.contact_model, contacts.ViscoElasticContacts)
 
-            # It is not yet clear how to pass the time step to this stage.
-            # A possibility is to restrict the integrator to only forward Euler
-            # and store the Δt inside the model.
-            module = jaxsim.rbda.contacts.visco_elastic.step.__module__
-            name = jaxsim.rbda.contacts.visco_elastic.step.__name__
-            msg = "You need to use the custom '{}.{}' function with this contact model."
-            jaxsim.exceptions.raise_runtime_error_if(
-                condition=True, msg=msg.format(module, name)
-            )
-
-            # Compute the 6D force expressed in the inertial frame and applied to each
-            # collidable point.
-            W_f_Ci, (W_f̿_Ci, m_tf) = model.contact_model.compute_contact_forces(
-                model=model,
-                data=data,
-                dt=None,  # TODO
-                link_forces=link_forces,
-                joint_force_references=joint_force_references,
-            )
-
-            aux_data = dict(W_f_avg2_C=W_f̿_Ci, m_tf=m_tf)
+            kwargs_contact_model = common_kwargs | dict(dt=model.time_step) | kwargs
 
         case _:
-            raise ValueError(f"Invalid contact model {model.contact_model}")
+            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 each collidable point.
     # In inertial-fixed representation, the computation of these transforms
     # is not necessary and the conversion below becomes a no-op.
-    W_H_Ci = (
+    W_H_C = (
         js.contact.transforms(model=model, data=data)
         if data.velocity_representation is not VelRepr.Inertial
         else jnp.zeros(
@@ -253,7 +218,7 @@ def collidable_point_dynamics(
             transform=W_H_C,
             is_force=True,
         )
-    )(W_f_Ci, W_H_Ci)
+    )(W_f_C, W_H_C)
 
     return f_Ci, aux_data
 
@@ -392,11 +357,13 @@ def estimate_good_contact_parameters(
                 max_penetration=max_δ,
                 number_of_active_collidable_points_steady_state=nc,
                 damping_ratio=damping_ratio,
-                **dict(
-                    p=model.contact_model.parameters.p,
-                    q=model.contact_model.parameters.q,
-                )
-                | kwargs,
+                **(
+                    dict(
+                        p=model.contact_model.parameters.p,
+                        q=model.contact_model.parameters.q,
+                    )
+                    | kwargs
+                ),
             )
 
         case contacts.ViscoElasticContacts():
@@ -410,11 +377,13 @@ def estimate_good_contact_parameters(
                     max_penetration=max_δ,
                     number_of_active_collidable_points_steady_state=nc,
                     damping_ratio=damping_ratio,
-                    **dict(
-                        p=model.contact_model.parameters.p,
-                        q=model.contact_model.parameters.q,
-                    )
-                    | kwargs,
+                    **(
+                        dict(
+                            p=model.contact_model.parameters.p,
+                            q=model.contact_model.parameters.q,
+                        )
+                        | kwargs
+                    ),
                 )
             )
 
@@ -427,11 +396,13 @@ def estimate_good_contact_parameters(
 
             parameters = contacts.RigidContactsParams.build(
                 mu=static_friction_coefficient,
-                **dict(
-                    K=K,
-                    D=2 * jnp.sqrt(K),
-                )
-                | kwargs,
+                **(
+                    dict(
+                        K=K,
+                        D=2 * jnp.sqrt(K),
+                    )
+                    | kwargs
+                ),
             )
 
         case contacts.RelaxedRigidContacts():

{jaxsim-0.4.3.dev229 → jaxsim-0.4.3.dev242}/src/jaxsim/api/model.py

@@ -1770,8 +1770,10 @@ def link_bias_accelerations(
 def link_contact_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 contact forces of all links of the model.
@@ -1784,6 +1786,7 @@ def link_contact_forces(
             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:
         A `(nL, 6)` array containing the stacked 6D contact forces of the links,
@@ -1820,47 +1823,16 @@ def link_contact_forces(
         joint_force_references=joint_force_references,
     )
 
-    # Compute the 6D forces applied to each collidable point expressed in the
-    # inertial frame.
-    with (
-        data.switch_velocity_representation(VelRepr.Inertial),
-        input_references.switch_velocity_representation(VelRepr.Inertial),
-    ):
-        W_f_C = js.contact.collidable_point_forces(
-            model=model,
-            data=data,
-            link_forces=input_references.link_forces(),
-            joint_force_references=input_references.joint_force_references(),
-        )
-
-    # 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(
-        model.kin_dyn_parameters.contact_parameters.body, dtype=int
-    )
-
-    # 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
-
-    # Create a references object to store the link forces.
-    references = js.references.JaxSimModelReferences.build(
-        model=model, link_forces=W_f_L, velocity_representation=VelRepr.Inertial
+    # Compute the 6D forces applied to the links equivalent to the forces applied
+    # to the frames associated to the collidable points.
+    f_L, _ = model.contact_model.compute_link_contact_forces(
+        model=model,
+        data=data,
+        link_forces=input_references.link_forces(model=model, data=data),
+        joint_force_references=input_references.joint_force_references(),
+        **kwargs,
     )
 
-    # Use the references object to convert the link forces 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
 
 
@@ -1967,6 +1939,11 @@ def step(
     Returns:
         A tuple containing the new data of the model
         and the new state of the integrator.
+
+    Note:
+        In order to reduce the occurrences of frame conversions performed internally,
+        it is recommended to use inertial-fixed velocity representation. This can be
+        particularly useful for automatically differentiated logic.
     """
 
     # Extract the integrator kwargs.
@@ -1976,15 +1953,61 @@ def step(
     integrator_kwargs = kwargs.pop("integrator_kwargs", {})
     integrator_kwargs = kwargs | integrator_kwargs
 
-    integrator_state = integrator_state if integrator_state is not None else dict()
+    # Initialize the integrator state.
+    integrator_state_t0 = integrator_state if integrator_state is not None else dict()
 
     # Initialize the time-related variables.
     state_t0 = data.state
     t0 = jnp.array(t0, dtype=float)
     dt = jnp.array(dt if dt is not None else model.time_step).astype(float)
 
-    # Rename the integrator state.
-    integrator_state_t0 = integrator_state
+    # The visco-elastic contacts operate at best with their own integrator.
+    # They can be used with Euler-like integrators, paying the price of ignoring
+    # some of the benefits of continuous-time integration on the system position.
+    # Furthermore, the requirement to know the Δt used by the integrator is not
+    # compatible with high-order integrators, that use advanced RK stages to evaluate
+    # the dynamics at intermediate times.
+    module = jaxsim.rbda.contacts.visco_elastic.step.__module__
+    name = jaxsim.rbda.contacts.visco_elastic.step.__name__
+    msg = "You need to use the custom '{}.{}' function with this contact model."
+    jaxsim.exceptions.raise_runtime_error_if(
+        condition=(
+            isinstance(model.contact_model, jaxsim.rbda.contacts.ViscoElasticContacts)
+            & (
+                ~jnp.allclose(dt, model.time_step)
+                | ~isinstance(integrator, jaxsim.integrators.fixed_step.ForwardEuler)
+            )
+        ),
+        msg=msg.format(module, name),
+    )
+
+    # =================
+    # Phase 1: pre-step
+    # =================
+
+    # TODO: some contact models here may want to perform a dynamic filtering of
+    # the enabled collidable points.
+
+    # Build the references object.
+    # We assume that the link forces are expressed in the frame corresponding to the
+    # velocity representation of the data.
+    references = js.references.JaxSimModelReferences.build(
+        model=model,
+        data=data,
+        velocity_representation=data.velocity_representation,
+        link_forces=link_forces,
+        joint_force_references=joint_force_references,
+    )
+
+    # =============
+    # Phase 2: step
+    # =============
+
+    # Prepare the references to pass.
+    with references.switch_velocity_representation(data.velocity_representation):
+
+        f_L = references.link_forces(model=model, data=data)
+        τ_references = references.joint_force_references(model=model)
 
     # Step the dynamics forward.
     state_tf, integrator_state_tf = integrator.step(
@@ -1994,7 +2017,7 @@ def step(
         params=integrator_state_t0,
         # Always inject the current (model, data) pair into the system dynamics
         # considered by the integrator, and include the input variables represented
-        # by the pair (joint_force_references, link_forces).
+        # by the pair (f_L, τ_references).
         # Note that the wrapper of the system dynamics will override (state_x0, t0)
         # inside the passed data even if it is not strictly needed. This logic is
         # necessary to re-use the jit-compiled step function of compatible pytrees
@@ -2003,8 +2026,8 @@ def step(
             dict(
                 model=model,
                 data=data,
-                joint_force_references=joint_force_references,
-                link_forces=link_forces,
+                link_forces=f_L,
+                joint_force_references=τ_references,
             )
             | integrator_kwargs
         ),
@@ -2013,6 +2036,10 @@ def step(
     # Store the new state of the model.
     data_tf = data.replace(state=state_tf)
 
+    # ==================
+    # Phase 3: post-step
+    # ==================
+
     # Post process the simulation state, if needed.
     match model.contact_model:
 
@@ -2040,17 +2067,18 @@ def step(
                 msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
             )
 
+            W_p_C = js.contact.collidable_point_positions(model, data_tf)
+
+            # Compute the penetration depth of the collidable points.
+            δ, *_ = jax.vmap(
+                jaxsim.rbda.contacts.common.compute_penetration_data,
+                in_axes=(0, 0, None),
+            )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
+
             with data_tf.switch_velocity_representation(VelRepr.Mixed):
 
                 J_WC = js.contact.jacobian(model, data_tf)
                 M = js.model.free_floating_mass_matrix(model, data_tf)
-                W_p_C = js.contact.collidable_point_positions(model, data_tf)
-
-                # Compute the penetration depth of the collidable points.
-                δ, *_ = jax.vmap(
-                    jaxsim.rbda.contacts.common.compute_penetration_data,
-                    in_axes=(0, 0, None),
-                )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
 
                 # Compute the impact velocity.
                 # It may be discontinuous in case new contacts are made.
@@ -2063,13 +2091,13 @@ def step(
                     )
                 )
 
-                # Reset the generalized velocity.
-                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:])
+            # Reset the generalized velocity.
+            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
-            )
+    # Restore the input velocity representation.
+    data_tf = data_tf.replace(
+        velocity_representation=data.velocity_representation, validate=False
+    )
 
     return data_tf, integrator_state_tf

{jaxsim-0.4.3.dev229 → jaxsim-0.4.3.dev242}/src/jaxsim/api/ode.py

@@ -131,7 +131,7 @@ def system_velocity_dynamics(
 
     # 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(O_f_L).astype(float)
+    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.
@@ -139,66 +139,59 @@ def system_velocity_dynamics(
 
     if len(model.kin_dyn_parameters.contact_parameters.body) > 0:
 
-        # Note: the following code should be kept in sync with the function
-        # `jaxsim.api.model.link_contact_forces`. We cannot merge them since
-        # here we need to get also aux_data.
-
-        # Compute the 6D forces W_f ∈ ℝ^{n_c × 6} applied to each collidable point
-        # along with contact-specific auxiliary states.
         with (
             data.switch_velocity_representation(VelRepr.Inertial),
             references.switch_velocity_representation(VelRepr.Inertial),
         ):
-            W_f_Ci, aux_data = js.contact.collidable_point_dynamics(
+
+            # 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),
             )
 
-        # 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(
-            model.kin_dyn_parameters.contact_parameters.body, dtype=int
-        )
-
-        # 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.
-        mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
-            model.number_of_links()
-        )
-
-        W_f_Li_terrain = mask.T @ W_f_Ci
+            # 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
+    # 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_Li_terrain,
+            forces=W_f_L_terrain,
             additive=True,
         )
 
-        # Get the link forces in inertial representation
+        # Get the link forces in inertial-fixed representation.
         f_L_total = references.link_forces(model=model, data=data)
 
-        v̇_WB, s̈ = system_acceleration(
+        # 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 v̇_WB, s̈, aux_data
+    return W_v̇_WB, s̈, aux_data
 
 
 def system_acceleration(
@@ -390,17 +383,15 @@ def system_dynamics(
 
         case contacts.ViscoElasticContacts():
 
-            extended_ode_state["contacts_state"] = {
-                "tangential_deformation": jnp.zeros_like(
-                    data.state.extended["tangential_deformation"]
-                )
-            }
+            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}")
+            raise ValueError(f"Invalid contact model: {model.contact_model}")
 
     # Extract the velocities.
     W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(