jaxsim 0.4.3.dev181__py3-none-any.whl → 0.4.3.dev200__py3-none-any.whl

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -12,11 +12,10 @@ import optax
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim import logging
-from jaxsim.api.common import VelRepr
-from jaxsim.math import Adjoint
+from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
 from jaxsim.terrain.terrain import FlatTerrain, Terrain
 
-from .common import ContactModel, ContactsParams
+from . import common
 
 try:
     from typing import Self
@@ -25,7 +24,7 @@ except ImportError:
 
 
 @jax_dataclasses.pytree_dataclass
-class RelaxedRigidContactsParams(ContactsParams):
+class RelaxedRigidContactsParams(common.ContactsParams):
     """Parameters of the relaxed rigid contacts model."""
 
     # Time constant
@@ -116,14 +115,24 @@ class RelaxedRigidContactsParams(ContactsParams):
     ) -> Self:
         """Create a `RelaxedRigidContactsParams` instance"""
 
+        def default(name: str):
+            return cls.__dataclass_fields__[name].default_factory()
+
         return cls(
-            **{
-                field: jnp.array(locals().get(field, default), dtype=default.dtype)
-                for field, default in map(
-                    lambda f: (f, cls.__dataclass_fields__[f].default),
-                    filter(lambda f: f != "__mutability__", cls.__dataclass_fields__),
-                )
-            }
+            time_constant=jnp.array(
+                time_constant or default("time_constant"), dtype=float
+            ),
+            damping_coefficient=jnp.array(
+                damping_coefficient or default("damping_coefficient"), dtype=float
+            ),
+            d_min=jnp.array(d_min or default("d_min"), dtype=float),
+            d_max=jnp.array(d_max or default("d_max"), dtype=float),
+            width=jnp.array(width or default("width"), dtype=float),
+            midpoint=jnp.array(midpoint or default("midpoint"), dtype=float),
+            power=jnp.array(power or default("power"), dtype=float),
+            stiffness=jnp.array(stiffness or default("stiffness"), dtype=float),
+            damping=jnp.array(damping or default("damping"), dtype=float),
+            mu=jnp.array(mu or default("mu"), dtype=float),
         )
 
     def valid(self) -> jtp.BoolLike:
@@ -142,7 +151,7 @@ class RelaxedRigidContactsParams(ContactsParams):
 
 
 @jax_dataclasses.pytree_dataclass
-class RelaxedRigidContacts(ContactModel):
+class RelaxedRigidContacts(common.ContactModel):
     """Relaxed rigid contacts model."""
 
     parameters: RelaxedRigidContactsParams = dataclasses.field(
@@ -229,7 +238,7 @@ class RelaxedRigidContacts(ContactModel):
         *,
         link_forces: jtp.MatrixLike | None = None,
         joint_force_references: jtp.VectorLike | None = None,
-    ) -> tuple[jtp.Vector, tuple[Any, ...]]:
+    ) -> tuple[jtp.Matrix, tuple]:
         """
         Compute the contact forces.
 
@@ -243,22 +252,23 @@ class RelaxedRigidContacts(ContactModel):
                 Optional `(n_joints,)` vector of joint forces.
 
         Returns:
-            A tuple containing the contact forces.
+            A tuple containing as first element the computed contact forces.
         """
 
         # Initialize the model and data this contact model is operating on.
         # This will raise an exception if either the contact model or the
         # contact parameters are not compatible.
         model, data = self.initialize_model_and_data(model=model, data=data)
+        assert isinstance(data.contacts_params, RelaxedRigidContactsParams)
 
-        link_forces = (
-            link_forces
+        link_forces = jnp.atleast_2d(
+            jnp.array(link_forces, dtype=float).squeeze()
             if link_forces is not None
             else jnp.zeros((model.number_of_links(), 6))
         )
 
-        joint_force_references = (
-            joint_force_references
+        joint_force_references = jnp.atleast_1d(
+            jnp.array(joint_force_references, dtype=float).squeeze()
             if joint_force_references is not None
             else jnp.zeros(model.number_of_joints())
         )
@@ -271,10 +281,10 @@ class RelaxedRigidContacts(ContactModel):
             joint_force_references=joint_force_references,
         )
 
-        def _detect_contact(x: jtp.Array, y: jtp.Array, z: jtp.Array) -> jtp.Array:
+        def detect_contact(x: jtp.Array, y: jtp.Array, z: jtp.Array) -> jtp.Array:
             x, y, z = jax.tree.map(jnp.squeeze, (x, y, z))
 
-            n̂ = self.terrain.normal(x=x, y=y).squeeze()
+            n̂ = model.terrain.normal(x=x, y=y).squeeze()
             h = jnp.array([0, 0, z - model.terrain.height(x=x, y=y)])
 
             return jnp.dot(h, n̂)
@@ -286,19 +296,19 @@ class RelaxedRigidContacts(ContactModel):
         )
 
         # Compute the activation state of the collidable points
-        δ = jax.vmap(_detect_contact)(*position.T)
+        δ = jax.vmap(detect_contact)(*position.T)
+
+        # Compute the transforms of the implicit frames corresponding to the
+        # collidable points.
+        W_H_C = js.contact.transforms(model=model, data=data)
 
         with (
             references.switch_velocity_representation(VelRepr.Mixed),
             data.switch_velocity_representation(VelRepr.Mixed),
         ):
-            M = js.model.free_floating_mass_matrix(model=model, data=data)
-            Jl_WC = jnp.vstack(
-                jax.vmap(lambda J, height: J * (height < 0))(
-                    js.contact.jacobian(model=model, data=data)[:, :3, :], δ
-                )
-            )
-            W_H_C = js.contact.transforms(model=model, data=data)
+
+            BW_ν = data.generalized_velocity()
+
             BW_ν̇_free = jnp.hstack(
                 js.ode.system_acceleration(
                     model=model,
@@ -309,20 +319,31 @@ class RelaxedRigidContacts(ContactModel):
                     ),
                 )
             )
-            BW_ν = data.generalized_velocity()
+
+            M = js.model.free_floating_mass_matrix(model=model, data=data)
+
+            Jl_WC = jnp.vstack(
+                jax.vmap(lambda J, height: J * (height < 0))(
+                    js.contact.jacobian(model=model, data=data)[:, :3, :], δ
+                )
+            )
+
             J̇_WC = jnp.vstack(
                 jax.vmap(lambda J̇, height: J̇ * (height < 0))(
                     js.contact.jacobian_derivative(model=model, data=data)[:, :3], δ
                 ),
             )
 
-            a_ref, R, K, D = self._regularizers(
-                model=model,
-                penetration=δ,
-                velocity=velocity,
-                parameters=self.parameters,
-            )
+        # Compute the regularization terms.
+        a_ref, R, K, D = self._regularizers(
+            model=model,
+            penetration=δ,
+            velocity=velocity,
+            parameters=data.contacts_params,
+        )
 
+        # Compute the Delassus matrix and the free mixed linear acceleration of
+        # the collidable points.
         G = Jl_WC @ jnp.linalg.lstsq(M, Jl_WC.T)[0]
         CW_al_free_WC = Jl_WC @ BW_ν̇_free + J̇_WC @ BW_ν
 
@@ -330,26 +351,40 @@ class RelaxedRigidContacts(ContactModel):
         A = G + R
         b = CW_al_free_WC - a_ref
 
+        # Create the objective function to minimize as a lambda computing the cost
+        # from the optimized variables x.
         objective = lambda x, A, b: jnp.sum(jnp.square(A @ x + b))
 
+        # ========================================
+        # Helper function to run the L-BFGS solver
+        # ========================================
+
         def run_optimization(
-            init_params: jtp.Array,
+            init_params: jtp.Vector,
             fun: Callable,
-            opt: optax.GradientTransformation,
-            maxiter: jtp.Int,
-            tol: jtp.Float,
-            **kwargs,
-        ):
+            opt: optax.GradientTransformationExtraArgs,
+            maxiter: int,
+            tol: float,
+        ) -> tuple[jtp.Vector, optax.OptState]:
+
+            # Get the function to compute the loss and the gradient w.r.t. its inputs.
             value_and_grad_fn = optax.value_and_grad_from_state(fun)
 
-            def step(carry):
+            # Initialize the carry of the following loop.
+            OptimizationCarry = tuple[jtp.Vector, optax.OptState]
+            init_carry: OptimizationCarry = (init_params, opt.init(params=init_params))
+
+            def step(carry: OptimizationCarry) -> OptimizationCarry:
+
                 params, state = carry
+
                 value, grad = value_and_grad_fn(
                     params,
                     state=state,
                     A=A,
                     b=b,
                 )
+
                 updates, state = opt.update(
                     updates=grad,
                     state=state,
@@ -360,22 +395,32 @@ class RelaxedRigidContacts(ContactModel):
                     A=A,
                     b=b,
                 )
+
                 params = optax.apply_updates(params, updates)
+
                 return params, state
 
-            def continuing_criterion(carry):
+            def continuing_criterion(carry: OptimizationCarry) -> jtp.Bool:
+
                 _, state = carry
+
                 iter_num = optax.tree_utils.tree_get(state, "count")
                 grad = optax.tree_utils.tree_get(state, "grad")
                 err = optax.tree_utils.tree_l2_norm(grad)
+
                 return (iter_num == 0) | ((iter_num < maxiter) & (err >= tol))
 
-            init_carry = (init_params, opt.init(init_params))
             final_params, final_state = jax.lax.while_loop(
                 continuing_criterion, step, init_carry
             )
+
             return final_params, final_state
 
+        # ======================================
+        # Compute the contact forces with L-BFGS
+        # ======================================
+
+        # Initialize the optimized forces with a linear Hunt/Crossley model.
         init_params = (
             K[:, jnp.newaxis] * jnp.zeros_like(position).at[:, 2].set(δ)
             + D[:, jnp.newaxis] * velocity
@@ -390,28 +435,30 @@ class RelaxedRigidContacts(ContactModel):
         maxiter = solver_options.pop("maxiter")
 
         # Compute the 3D linear force in C[W] frame.
-        CW_f_Ci, _ = run_optimization(
+        solution, _ = run_optimization(
             init_params=init_params,
-            A=A,
-            b=b,
-            maxiter=maxiter,
-            opt=optax.lbfgs(**solver_options),
             fun=objective,
+            opt=optax.lbfgs(**solver_options),
             tol=tol,
+            maxiter=maxiter,
         )
 
-        CW_f_Ci = CW_f_Ci.reshape((-1, 3))
-
-        def mixed_to_inertial(W_H_C: jax.Array, CW_fl: jax.Array) -> jax.Array:
-            W_Xf_CW = Adjoint.from_transform(
-                W_H_C.at[0:3, 0:3].set(jnp.eye(3)),
-                inverse=True,
-            ).T
-            return W_Xf_CW @ jnp.hstack([CW_fl, jnp.zeros(3)])
-
-        W_f_C = jax.vmap(mixed_to_inertial)(W_H_C, CW_f_Ci)
+        # Reshape the optimized solution to be a matrix of 3D contact forces.
+        CW_fl_C = solution.reshape(-1, 3)
+
+        # Convert the contact forces from mixed to inertial-fixed representation.
+        W_f_C = jax.vmap(
+            lambda CW_fl_C, W_H_C: (
+                ModelDataWithVelocityRepresentation.other_representation_to_inertial(
+                    array=jnp.zeros(6).at[0:3].set(CW_fl_C),
+                    transform=W_H_C,
+                    other_representation=VelRepr.Mixed,
+                    is_force=True,
+                )
+            ),
+        )(CW_fl_C, W_H_C)
 
-        return W_f_C, (None,)
+        return W_f_C, ()
 
     @staticmethod
     def _regularizers(
@@ -433,13 +480,28 @@ class RelaxedRigidContacts(ContactModel):
             A tuple containing the reference acceleration, the regularization matrix, the stiffness, and the damping.
         """
 
-        Ω, ζ, ξ_min, ξ_max, width, mid, p, K, D, μ, *_ = jax_dataclasses.astuple(
-            parameters
+        # Extract the parameters of the contact model.
+        Ω, ζ, ξ_min, ξ_max, width, mid, p, K, D, μ = (
+            getattr(parameters, field)
+            for field in (
+                "time_constant",
+                "damping_coefficient",
+                "d_min",
+                "d_max",
+                "width",
+                "midpoint",
+                "power",
+                "stiffness",
+                "damping",
+                "mu",
+            )
         )
 
-        def _imp_aref(
-            penetration: jtp.Array,
-            velocity: jtp.Array,
+        # Compute the 6D inertia matrices of all links.
+        M_L = js.model.link_spatial_inertia_matrices(model=model)
+
+        def imp_aref(
+            penetration: jtp.Array, velocity: jtp.Array
         ) -> tuple[jtp.Array, jtp.Array]:
             """
             Calculates impedance and offset acceleration in constraint frame.
@@ -474,7 +536,7 @@ class RelaxedRigidContacts(ContactModel):
 
             return imp, a_ref, jnp.atleast_1d(K_f), jnp.atleast_1d(D_f)
 
-        def _compute_row(
+        def compute_row(
             *,
             link_idx: jtp.Float,
             penetration: jtp.Array,
@@ -482,7 +544,7 @@ class RelaxedRigidContacts(ContactModel):
         ) -> tuple[jtp.Array, jtp.Array]:
 
             # Compute the reference acceleration.
-            ξ, a_ref, K, D = _imp_aref(
+            ξ, a_ref, K, D = imp_aref(
                 penetration=penetration,
                 velocity=velocity,
             )
@@ -496,12 +558,10 @@ class RelaxedRigidContacts(ContactModel):
 
             return jax.tree.map(lambda x: x * (penetration < 0), (a_ref, R, K, D))
 
-        M_L = js.model.link_spatial_inertia_matrices(model=model)
-
         a_ref, R, K, D = jax.tree.map(
-            jnp.concatenate,
-            (
-                *jax.vmap(_compute_row)(
+            f=jnp.concatenate,
+            tree=(
+                *jax.vmap(compute_row)(
                     link_idx=jnp.array(
                         model.kin_dyn_parameters.contact_parameters.body
                     ),
@@ -510,4 +570,5 @@ class RelaxedRigidContacts(ContactModel):
                 ),
             ),
         )
+
         return a_ref, jnp.diag(R), K, D