jaxsim 0.4.3.dev68__py3-none-any.whl → 0.4.3.dev70__py3-none-any.whl

jaxsim/parsers/rod/parser.py

@@ -223,7 +223,7 @@ def extract_model_data(
             child=links_dict[j.child],
             jtype=utils.joint_to_joint_type(joint=j),
             axis=(
-                np.array(j.axis.xyz.xyz)
+                np.array(j.axis.xyz.xyz, dtype=float)
                 if j.axis is not None
                 and j.axis.xyz is not None
                 and j.axis.xyz.xyz is not None
@@ -232,39 +232,43 @@ def extract_model_data(
             pose=j.pose.transform() if j.pose is not None else np.eye(4),
             initial_position=0.0,
             position_limit=(
-                (
-                    float(j.axis.limit.lower)
-                    if j.axis is not None and j.axis.limit is not None
-                    else np.finfo(float).min
+                float(
+                    j.axis.limit.lower
+                    if j.axis is not None
+                    and j.axis.limit is not None
+                    and j.axis.limit.lower is not None
+                    else jnp.finfo(float).min
                 ),
-                (
-                    float(j.axis.limit.upper)
-                    if j.axis is not None and j.axis.limit is not None
-                    else np.finfo(float).max
+                float(
+                    j.axis.limit.upper
+                    if j.axis is not None
+                    and j.axis.limit is not None
+                    and j.axis.limit.upper is not None
+                    else jnp.finfo(float).max
                 ),
             ),
-            friction_static=(
+            friction_static=float(
                 j.axis.dynamics.friction
                 if j.axis is not None
                 and j.axis.dynamics is not None
                 and j.axis.dynamics.friction is not None
                 else 0.0
             ),
-            friction_viscous=(
+            friction_viscous=float(
                 j.axis.dynamics.damping
                 if j.axis is not None
                 and j.axis.dynamics is not None
                 and j.axis.dynamics.damping is not None
                 else 0.0
             ),
-            position_limit_damper=(
+            position_limit_damper=float(
                 j.axis.limit.dissipation
                 if j.axis is not None
                 and j.axis.limit is not None
                 and j.axis.limit.dissipation is not None
                 else 0.0
             ),
-            position_limit_spring=(
+            position_limit_spring=float(
                 j.axis.limit.stiffness
                 if j.axis is not None
                 and j.axis.limit is not None
@@ -273,7 +277,7 @@ def extract_model_data(
             ),
         )
         for j in sdf_model.joints()
-        if j.type in {"revolute", "prismatic", "fixed"}
+        if j.type in {"revolute", "continuous", "prismatic", "fixed"}
         and j.parent != "world"
         and j.child in links_dict.keys()
     ]

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -0,0 +1,384 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import jax_dataclasses
+import jaxopt
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.api.common import VelRepr
+from jaxsim.math import Adjoint
+from jaxsim.terrain.terrain import FlatTerrain, Terrain
+
+from .common import ContactModel, ContactsParams, ContactsState
+
+
+@jax_dataclasses.pytree_dataclass
+class RelaxedRigidContactsParams(ContactsParams):
+    """Parameters of the relaxed rigid contacts model."""
+
+    # Time constant
+    time_constant: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.01, dtype=float)
+    )
+
+    # Adimensional damping coefficient
+    damping_coefficient: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(1.0, dtype=float)
+    )
+
+    # Minimum impedance
+    d_min: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.9, dtype=float)
+    )
+
+    # Maximum impedance
+    d_max: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.95, dtype=float)
+    )
+
+    # Width
+    width: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.0001, dtype=float)
+    )
+
+    # Midpoint
+    midpoint: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.1, dtype=float)
+    )
+
+    # Power exponent
+    power: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(1.0, dtype=float)
+    )
+
+    # Stiffness
+    stiffness: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.0, dtype=float)
+    )
+
+    # Damping
+    damping: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.0, dtype=float)
+    )
+
+    # Friction coefficient
+    mu: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(0.5, dtype=float)
+    )
+
+    # Maximum number of iterations
+    max_iterations: jtp.Int = dataclasses.field(
+        default_factory=lambda: jnp.array(50, dtype=int)
+    )
+
+    # Solver tolerance
+    tolerance: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(1e-6, dtype=float)
+    )
+
+    def __hash__(self) -> int:
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return hash(
+            (
+                HashedNumpyArray(self.time_constant),
+                HashedNumpyArray(self.damping_coefficient),
+                HashedNumpyArray(self.d_min),
+                HashedNumpyArray(self.d_max),
+                HashedNumpyArray(self.width),
+                HashedNumpyArray(self.midpoint),
+                HashedNumpyArray(self.power),
+                HashedNumpyArray(self.stiffness),
+                HashedNumpyArray(self.damping),
+                HashedNumpyArray(self.mu),
+                HashedNumpyArray(self.max_iterations),
+                HashedNumpyArray(self.tolerance),
+            )
+        )
+
+    def __eq__(self, other: RelaxedRigidContactsParams) -> bool:
+        return hash(self) == hash(other)
+
+    @classmethod
+    def build(
+        cls,
+        time_constant: jtp.FloatLike | None = None,
+        damping_coefficient: jtp.FloatLike | None = None,
+        d_min: jtp.FloatLike | None = None,
+        d_max: jtp.FloatLike | None = None,
+        width: jtp.FloatLike | None = None,
+        midpoint: jtp.FloatLike | None = None,
+        power: jtp.FloatLike | None = None,
+        stiffness: jtp.FloatLike | None = None,
+        damping: jtp.FloatLike | None = None,
+        mu: jtp.FloatLike | None = None,
+        max_iterations: jtp.IntLike | None = None,
+        tolerance: jtp.FloatLike | None = None,
+    ) -> RelaxedRigidContactsParams:
+        """Create a `RelaxedRigidContactsParams` instance"""
+
+        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__),
+                )
+            }
+        )
+
+    def valid(self) -> bool:
+        return bool(
+            jnp.all(self.time_constant >= 0.0)
+            and jnp.all(self.damping_coefficient > 0.0)
+            and jnp.all(self.d_min >= 0.0)
+            and jnp.all(self.d_max <= 1.0)
+            and jnp.all(self.d_min <= self.d_max)
+            and jnp.all(self.width >= 0.0)
+            and jnp.all(self.midpoint >= 0.0)
+            and jnp.all(self.power >= 0.0)
+            and jnp.all(self.mu >= 0.0)
+            and jnp.all(self.max_iterations > 0)
+            and jnp.all(self.tolerance > 0.0)
+        )
+
+
+@jax_dataclasses.pytree_dataclass
+class RelaxedRigidContactsState(ContactsState):
+    """Class storing the state of the relaxed rigid contacts model."""
+
+    def __eq__(self, other: RelaxedRigidContactsState) -> bool:
+        return hash(self) == hash(other)
+
+    @staticmethod
+    def build() -> RelaxedRigidContactsState:
+        """Create a `RelaxedRigidContactsState` instance"""
+
+        return RelaxedRigidContactsState()
+
+    @staticmethod
+    def zero(model: js.model.JaxSimModel) -> RelaxedRigidContactsState:
+        """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
+        return RelaxedRigidContactsState.build()
+
+    def valid(self, model: js.model.JaxSimModel) -> bool:
+        return True
+
+
+@jax_dataclasses.pytree_dataclass
+class RelaxedRigidContacts(ContactModel):
+    """Relaxed rigid contacts model."""
+
+    parameters: RelaxedRigidContactsParams = dataclasses.field(
+        default_factory=RelaxedRigidContactsParams
+    )
+
+    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
+        default_factory=FlatTerrain
+    )
+
+    def compute_contact_forces(
+        self,
+        position: jtp.Vector,
+        velocity: jtp.Vector,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        link_forces: jtp.MatrixLike | None = None,
+    ) -> tuple[jtp.Vector, tuple[Any, ...]]:
+
+        link_forces = (
+            link_forces
+            if link_forces is not None
+            else jnp.zeros((model.number_of_links(), 6))
+        )
+
+        references = js.references.JaxSimModelReferences.build(
+            model=model,
+            data=data,
+            velocity_representation=data.velocity_representation,
+            link_forces=link_forces,
+        )
+
+        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()
+            h = jnp.array([0, 0, z - model.terrain.height(x=x, y=y)])
+
+            return jnp.dot(h, n̂)
+
+        # Compute the activation state of the collidable points
+        δ = jax.vmap(_detect_contact)(*position.T)
+
+        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_ν̇_free = jnp.hstack(
+                js.ode.system_acceleration(
+                    model=model,
+                    data=data,
+                    link_forces=references.link_forces(model=model, data=data),
+                )
+            )
+            BW_ν = data.generalized_velocity()
+            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,
+            )
+
+        G = Jl_WC @ jnp.linalg.lstsq(M, Jl_WC.T)[0]
+        CW_al_free_WC = Jl_WC @ BW_ν̇_free + J̇_WC @ BW_ν
+
+        # Calculate quantities for the linear optimization problem.
+        A = G + R
+        b = CW_al_free_WC - a_ref
+
+        objective = lambda x: jnp.sum(jnp.square(A @ x + b))
+
+        # Compute the 3D linear force in C[W] frame
+        opt = jaxopt.LBFGS(
+            fun=objective,
+            maxiter=self.parameters.max_iterations,
+            tol=self.parameters.tolerance,
+            maxls=30,
+            history_size=10,
+            max_stepsize=100.0,
+        )
+
+        init_params = (
+            K[:, jnp.newaxis] * jnp.zeros_like(position).at[:, 2].set(δ)
+            + D[:, jnp.newaxis] * velocity
+        ).flatten()
+
+        CW_f_Ci = opt.run(init_params=init_params).params.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)
+
+        return W_f_C, (None,)
+
+    @staticmethod
+    def _regularizers(
+        model: js.model.JaxSimModel,
+        penetration: jtp.Array,
+        velocity: jtp.Array,
+        parameters: RelaxedRigidContactsParams,
+    ) -> tuple:
+        """
+        Compute the contact jacobian and the reference acceleration.
+
+        Args:
+            model: The jaxsim model.
+            penetration: The penetration of the collidable points.
+            velocity: The velocity of the collidable points.
+            parameters: The parameters of the relaxed rigid contacts model.
+
+        Returns:
+            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
+        )
+
+        def _imp_aref(
+            penetration: jtp.Array,
+            velocity: jtp.Array,
+        ) -> tuple[jtp.Array, jtp.Array]:
+            """
+            Calculates impedance and offset acceleration in constraint frame.
+
+            Args:
+                penetration: penetration in constraint frame
+                velocity: velocity in constraint frame
+
+            Returns:
+                a_ref: offset acceleration in constraint frame
+                R: regularization matrix
+                K: computed stiffness
+                D: computed damping
+            """
+            position = jnp.zeros(shape=(3,)).at[2].set(penetration)
+
+            imp_x = jnp.abs(position) / width
+            imp_a = (1.0 / jnp.power(mid, p - 1)) * jnp.power(imp_x, p)
+
+            imp_b = 1 - (1.0 / jnp.power(1 - mid, p - 1)) * jnp.power(1 - imp_x, p)
+
+            imp_y = jnp.where(imp_x < mid, imp_a, imp_b)
+
+            imp = jnp.clip(ξ_min + imp_y * (ξ_max - ξ_min), ξ_min, ξ_max)
+            imp = jnp.atleast_1d(jnp.where(imp_x > 1.0, ξ_max, imp))
+
+            # When passing negative values, K and D represent a spring and damper, respectively.
+            K_f = jnp.where(K < 0, -K / ξ_max**2, 1 / (ξ_max * Ω * ζ) ** 2)
+            D_f = jnp.where(D < 0, -D / ξ_max, 2 / (ξ_max * Ω))
+
+            a_ref = -jnp.atleast_1d(D_f * velocity + K_f * imp * position)
+
+            return imp, a_ref, jnp.atleast_1d(K_f), jnp.atleast_1d(D_f)
+
+        def _compute_row(
+            *,
+            link_idx: jtp.Float,
+            penetration: jtp.Array,
+            velocity: jtp.Array,
+        ) -> tuple[jtp.Array, jtp.Array]:
+
+            # Compute the reference acceleration.
+            ξ, a_ref, K, D = _imp_aref(
+                penetration=penetration,
+                velocity=velocity,
+            )
+
+            # Compute the regularization terms.
+            R = (
+                (2 * μ**2 * (1 - ξ) / (ξ + 1e-12))
+                * (1 + μ**2)
+                @ jnp.linalg.inv(M_L[link_idx, :3, :3])
+            )
+
+            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)(
+                    link_idx=jnp.array(
+                        model.kin_dyn_parameters.contact_parameters.body
+                    ),
+                    penetration=penetration,
+                    velocity=velocity,
+                ),
+            ),
+        )
+        return a_ref, jnp.diag(R), K, D

jaxsim/rbda/contacts/rigid.py

@@ -9,7 +9,6 @@ import jax_dataclasses
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim import math
 from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
 from jaxsim.terrain import FlatTerrain, Terrain
 
@@ -272,9 +271,17 @@ class RigidContacts(ContactModel):
             link_forces=link_forces,
         )
 
-        with references.switch_velocity_representation(VelRepr.Mixed):
-            BW_ν̇_free = RigidContacts._compute_mixed_nu_dot_free(
-                model, data, references=references
+        with (
+            references.switch_velocity_representation(VelRepr.Mixed),
+            data.switch_velocity_representation(VelRepr.Mixed),
+        ):
+            BW_ν̇_free = jnp.hstack(
+                js.ode.system_acceleration(
+                    model=model,
+                    data=data,
+                    joint_forces=references.joint_force_references(model=model),
+                    link_forces=references.link_forces(model=model, data=data),
+                )
             )
 
         free_contact_acc = RigidContacts._linear_acceleration_of_collidable_points(
@@ -380,43 +387,6 @@ class RigidContacts(ContactModel):
         n_constraints = 6 * n_collidable_points
         return jnp.zeros(shape=(n_constraints,))
 
-    @staticmethod
-    def _compute_mixed_nu_dot_free(
-        model: js.model.JaxSimModel,
-        data: js.data.JaxSimModelData,
-        references: js.references.JaxSimModelReferences | None = None,
-    ) -> jtp.Array:
-        references = (
-            references
-            if references is not None
-            else js.references.JaxSimModelReferences.zero(model=model, data=data)
-        )
-
-        with (
-            data.switch_velocity_representation(VelRepr.Mixed),
-            references.switch_velocity_representation(VelRepr.Mixed),
-        ):
-            BW_v_WB = data.base_velocity()
-            W_ṗ_B, W_ω_WB = jnp.split(BW_v_WB, 2)
-            W_v̇_WB, s̈ = js.ode.system_acceleration(
-                model=model,
-                data=data,
-                joint_forces=references.joint_force_references(model=model),
-                link_forces=references.link_forces(model=model, data=data),
-            )
-
-        # Convert the inertial-fixed base acceleration to a mixed base acceleration.
-        W_H_B = data.base_transform()
-        W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
-        BW_X_W = math.Adjoint.from_transform(W_H_BW, inverse=True)
-        term1 = BW_X_W @ W_v̇_WB
-        term2 = jnp.zeros(6).at[0:3].set(jnp.cross(W_ṗ_B, W_ω_WB))
-        BW_v̇_WB = term1 - term2
-
-        BW_ν̇ = jnp.hstack([BW_v̇_WB, s̈])
-
-        return BW_ν̇
-
     @staticmethod
     def _linear_acceleration_of_collidable_points(
         model: js.model.JaxSimModel,

jaxsim/terrain/terrain.py

@@ -46,66 +46,82 @@ class Terrain(abc.ABC):
 @jax_dataclasses.pytree_dataclass
 class FlatTerrain(Terrain):
 
-    z: float = dataclasses.field(default=0.0, kw_only=True)
+    _height: float = dataclasses.field(default=0.0, kw_only=True)
 
     @staticmethod
     def build(height: jtp.FloatLike) -> FlatTerrain:
 
-        return FlatTerrain(z=float(height))
+        return FlatTerrain(_height=float(height))
 
     def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
 
-        return jnp.array(self.z, dtype=float)
+        return jnp.array(self._height, dtype=float)
+
+    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
+
+        return jnp.array([0.0, 0.0, 1.0], dtype=float)
 
     def __hash__(self) -> int:
 
-        return hash(self.z)
+        return hash(self._height)
 
     def __eq__(self, other: FlatTerrain) -> bool:
 
         if not isinstance(other, FlatTerrain):
             return False
 
-        return self.z == other.z
+        return self._height == other._height
 
 
 @jax_dataclasses.pytree_dataclass
 class PlaneTerrain(FlatTerrain):
 
-    plane_normal: tuple[float, float, float] = jax_dataclasses.field(
+    _normal: tuple[float, float, float] = jax_dataclasses.field(
         default=(0.0, 0.0, 1.0), kw_only=True
     )
 
     @staticmethod
-    def build(
-        plane_normal: jtp.VectorLike, plane_height_over_origin: jtp.FloatLike = 0.0
-    ) -> PlaneTerrain:
+    def build(height: jtp.FloatLike = 0.0, *, normal: jtp.VectorLike) -> PlaneTerrain:
         """
         Create a PlaneTerrain instance with a specified plane normal vector.
 
         Args:
-            plane_normal: The normal vector of the terrain plane.
-            plane_height_over_origin: The height of the plane over the origin.
+            normal: The normal vector of the terrain plane.
+            height: The height of the plane over the origin.
 
         Returns:
             PlaneTerrain: A PlaneTerrain instance.
         """
 
-        plane_normal = jnp.array(plane_normal, dtype=float)
-        plane_height_over_origin = jnp.array(plane_height_over_origin, dtype=float)
+        normal = jnp.array(normal, dtype=float)
+        height = jnp.array(height, dtype=float)
 
-        if plane_normal.shape != (3,):
+        if normal.shape != (3,):
             msg = "Expected a 3D vector for the plane normal, got '{}'."
-            raise ValueError(msg.format(plane_normal.shape))
+            raise ValueError(msg.format(normal.shape))
 
         # Make sure that the plane normal is a unit vector.
-        plane_normal = plane_normal / jnp.linalg.norm(plane_normal)
+        normal = normal / jnp.linalg.norm(normal)
 
         return PlaneTerrain(
-            z=float(plane_height_over_origin),
-            plane_normal=tuple(plane_normal.tolist()),
+            _height=height.item(),
+            _normal=tuple(normal.tolist()),
         )
 
+    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
+        """
+        Compute the normal vector of the terrain at a specific (x, y) location.
+
+        Args:
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
+
+        Returns:
+            The normal vector of the terrain surface at the specified location.
+        """
+
+        return jnp.array(self._normal, dtype=float)
+
     def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
         """
         Compute the height of the terrain at a specific (x, y) location on a plane.
@@ -123,10 +139,10 @@ class PlaneTerrain(FlatTerrain):
         # The height over the origin: -D/C
 
         # Get the plane equation coefficients from the terrain normal.
-        A, B, C = self.plane_normal
+        A, B, C = self._normal
 
         # Compute the final coefficient D considering the terrain height.
-        D = -C * self.z
+        D = -C * self._height
 
         # Invert the plane equation to get the height at the given (x, y) coordinates.
         return jnp.array(-(A * x + B * y + D) / C).astype(float)
@@ -137,9 +153,9 @@ class PlaneTerrain(FlatTerrain):
 
         return hash(
             (
-                hash(self.z),
+                hash(self._height),
                 HashedNumpyArray.hash_of_array(
-                    array=jnp.array(self.plane_normal, dtype=float)
+                    array=jnp.array(self._normal, dtype=float)
                 ),
             )
         )
@@ -150,10 +166,10 @@ class PlaneTerrain(FlatTerrain):
             return False
 
         if not (
-            np.allclose(self.z, other.z)
+            np.allclose(self._height, other._height)
             and np.allclose(
-                np.array(self.plane_normal, dtype=float),
-                np.array(other.plane_normal, dtype=float),
+                np.array(self._normal, dtype=float),
+                np.array(other._normal, dtype=float),
             )
         ):
             return False

jaxsim/typing.py

@@ -16,7 +16,7 @@ Int = Scalar
 Bool = Scalar
 Float = Scalar
 
-PyTree = (
+PyTree: object = (
     dict[Hashable, TypeVar("PyTree")]
     | list[TypeVar("PyTree")]
     | tuple[TypeVar("PyTree")]