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

jaxsim/physics/algos/soft_contacts.py

@@ -1,523 +0,0 @@
-from __future__ import annotations
-
-import dataclasses
-from typing import Tuple
-
-import jax
-import jax.flatten_util
-import jax.numpy as jnp
-import jax_dataclasses
-import numpy as np
-
-import jaxsim.physics.model.physics_model
-import jaxsim.typing as jtp
-from jaxsim.math.adjoint import Adjoint
-from jaxsim.math.skew import Skew
-from jaxsim.physics.algos.terrain import FlatTerrain, Terrain
-from jaxsim.physics.model.physics_model import PhysicsModel
-from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass
-
-from . import utils
-
-
-@jax_dataclasses.pytree_dataclass
-class SoftContactsState(JaxsimDataclass):
-    """
-    State of the soft contacts model.
-
-    Attributes:
-        tangential_deformation:
-            The tangential deformation of the material at each collidable point.
-    """
-
-    tangential_deformation: jtp.Matrix
-
-    @staticmethod
-    def build(
-        tangential_deformation: jtp.Matrix | None = None,
-        number_of_collidable_points: int | None = None,
-    ) -> SoftContactsState:
-        """"""
-
-        tangential_deformation = (
-            tangential_deformation
-            if tangential_deformation is not None
-            else jnp.zeros(shape=(3, number_of_collidable_points))
-        )
-
-        return SoftContactsState(
-            tangential_deformation=jnp.array(tangential_deformation, dtype=float)
-        )
-
-    @staticmethod
-    def build_from_physics_model(
-        tangential_deformation: jtp.Matrix | None = None,
-        physics_model: jaxsim.physics.model.physics_model.PhysicsModel | None = None,
-    ) -> SoftContactsState:
-        """"""
-
-        return SoftContactsState.build(
-            tangential_deformation=tangential_deformation,
-            number_of_collidable_points=len(physics_model.gc.body),
-        )
-
-    @staticmethod
-    def zero(
-        physics_model: jaxsim.physics.model.physics_model.PhysicsModel,
-    ) -> SoftContactsState:
-        """
-        Modify the SoftContactsState instance imposing zero tangential deformation.
-
-        Args:
-            physics_model: The physics model.
-
-        Returns:
-            A SoftContactsState instance with zero tangential deformation.
-        """
-
-        return SoftContactsState.build_from_physics_model(physics_model=physics_model)
-
-    def valid(
-        self, physics_model: jaxsim.physics.model.physics_model.PhysicsModel
-    ) -> bool:
-        """
-        Check if the soft contacts state has valid shape.
-
-        Args:
-            physics_model: The physics model.
-
-        Returns:
-            True if the state has a valid shape, otherwise False.
-        """
-
-        from jaxsim.simulation.utils import check_valid_shape
-
-        return check_valid_shape(
-            what="tangential_deformation",
-            shape=self.tangential_deformation.shape,
-            expected_shape=(3, len(physics_model.gc.body)),
-            valid=True,
-        )
-
-
-def collidable_points_pos_vel(
-    model: PhysicsModel,
-    q: jtp.Vector,
-    qd: jtp.Vector,
-    xfb: jtp.Vector | None = None,
-) -> Tuple[jtp.Matrix, jtp.Matrix]:
-    """
-    Compute the position and linear velocity of collidable points in the world frame.
-
-    Args:
-        model (PhysicsModel): The physics model.
-        q (jtp.Vector): The joint positions.
-        qd (jtp.Vector): The joint velocities.
-        xfb (jtp.Vector, optional): The floating base state. Defaults to None.
-
-    Returns:
-        Tuple[jtp.Matrix, jtp.Matrix]: A tuple containing the position and velocity of collidable points.
-    """
-
-    # Make sure that shape and size are correct
-    xfb, q, qd, _, _, _ = utils.process_inputs(physics_model=model, xfb=xfb, q=q, qd=qd)
-
-    # Initialize buffers of link transforms (W_X_i) and 6D inertial velocities (W_v_Wi)
-    W_X_i = jnp.zeros(shape=[model.NB, 6, 6])
-    W_v_Wi = jnp.zeros(shape=[model.NB, 6, 1])
-
-    # 6D transform of base velocity
-    W_X_0 = Adjoint.from_quaternion_and_translation(
-        quaternion=xfb[0:4], translation=xfb[4:7], normalize_quaternion=True
-    )
-    W_X_i = W_X_i.at[0].set(W_X_0)
-
-    # Store the 6D inertial velocity W_v_W0 of the base link
-    W_v_W0 = jnp.vstack(jnp.hstack([xfb[10:13], xfb[7:10]]))
-    W_v_Wi = W_v_Wi.at[0].set(W_v_W0)
-
-    # Compute useful resources from the model
-    S = model.motion_subspaces(q=q)
-
-    # Get the 6D transform between the parent link λi and the joint's predecessor frame
-    pre_X_λi = model.tree_transforms
-
-    # Compute the 6D transform of the joints (from predecessor to successor)
-    i_X_pre = model.joint_transforms(q=q)
-
-    # Parent array mapping: i -> λ(i).
-    # Exception: λ(0) must not be used, it's initialized to -1.
-    λ = model.parent_array()
-
-    # ====================
-    # Propagate kinematics
-    # ====================
-
-    PropagateTransformsCarry = Tuple[jtp.MatrixJax]
-    propagate_transforms_carry: PropagateTransformsCarry = (W_X_i,)
-
-    def propagate_transforms(
-        carry: PropagateTransformsCarry, i: jtp.Int
-    ) -> Tuple[PropagateTransformsCarry, None]:
-        # Unpack the carry
-        (W_X_i,) = carry
-
-        # We need the inverse transforms (from parent to child direction)
-        pre_Xi_i = Adjoint.inverse(i_X_pre[i])
-        λi_Xi_pre = Adjoint.inverse(pre_X_λi[i])
-
-        # Compute the parent to child 6D transform
-        λi_X_i = λi_Xi_pre @ pre_Xi_i
-
-        # Compute the world to child 6D transform
-        W_Xi_i = W_X_i[λ[i]] @ λi_X_i
-        W_X_i = W_X_i.at[i].set(W_Xi_i)
-
-        # Pack and return the carry
-        return (W_X_i,), None
-
-    (W_X_i,), _ = jax.lax.scan(
-        f=propagate_transforms,
-        init=propagate_transforms_carry,
-        xs=np.arange(start=1, stop=model.NB),
-    )
-
-    # ====================
-    # Propagate velocities
-    # ====================
-
-    PropagateVelocitiesCarry = Tuple[jtp.MatrixJax]
-    propagate_velocities_carry: PropagateVelocitiesCarry = (W_v_Wi,)
-
-    def propagate_velocities(
-        carry: PropagateVelocitiesCarry, j_vel_and_j_idx: jtp.VectorJax
-    ) -> Tuple[PropagateVelocitiesCarry, None]:
-        # Unpack the scanned data
-        qd_ii = j_vel_and_j_idx[0]
-        ii = jnp.array(j_vel_and_j_idx[1], dtype=int)
-
-        # Given a joint whose velocity is qd[ii], the index of its parent link is ii + 1
-        i = ii + 1
-
-        # Unpack the carry
-        (W_v_Wi,) = carry
-
-        # Propagate the 6D velocity
-        W_vi_Wi = W_v_Wi[λ[i]] + W_X_i[i] @ (S[i] * qd_ii)
-        W_v_Wi = W_v_Wi.at[i].set(W_vi_Wi)
-
-        # Pack and return the carry
-        return (W_v_Wi,), None
-
-    (W_v_Wi,), _ = jax.lax.scan(
-        f=propagate_velocities,
-        init=propagate_velocities_carry,
-        xs=jnp.vstack([qd, jnp.arange(start=0, stop=qd.size)]).T,
-    )
-
-    # ==================================================
-    # Compute position and velocity of collidable points
-    # ==================================================
-
-    def process_point_kinematics(
-        Li_p_C: jtp.VectorJax, parent_body: jtp.Int
-    ) -> Tuple[jtp.VectorJax, jtp.VectorJax]:
-        # Compute the position of the collidable point
-        W_p_Ci = (
-            Adjoint.to_transform(adjoint=W_X_i[parent_body]) @ jnp.hstack([Li_p_C, 1])
-        )[0:3]
-
-        # Compute the linear part of the mixed velocity Ci[W]_v_{W,Ci}
-        CW_vl_WCi = (
-            jnp.block([jnp.eye(3), -Skew.wedge(vector=W_p_Ci).squeeze()])
-            @ W_v_Wi[parent_body].squeeze()
-        )
-
-        return W_p_Ci, CW_vl_WCi
-
-    # Process all the collidable points in parallel
-    W_p_Ci, CW_v_WC = jax.vmap(process_point_kinematics)(
-        model.gc.point.T, np.array(model.gc.body, dtype=int)
-    )
-
-    return W_p_Ci.transpose(), CW_v_WC.transpose()
-
-
-@jax_dataclasses.pytree_dataclass
-class SoftContactsParams:
-    """Parameters of the soft contacts model."""
-
-    K: float = dataclasses.field(default_factory=lambda: jnp.array(1e6, dtype=float))
-    D: float = dataclasses.field(default_factory=lambda: jnp.array(2000, dtype=float))
-    mu: float = dataclasses.field(default_factory=lambda: jnp.array(0.5, dtype=float))
-
-    @staticmethod
-    def build(
-        K: jtp.FloatLike = 1e6, D: jtp.FloatLike = 2_000, mu: jtp.FloatLike = 0.5
-    ) -> SoftContactsParams:
-        """
-        Create a SoftContactsParams instance with specified parameters.
-
-        Args:
-            K (float, optional): The stiffness parameter. Defaults to 1e6.
-            D (float, optional): The damping parameter. Defaults to 2000.
-            mu (float, optional): The friction coefficient. Defaults to 0.5.
-
-        Returns:
-            SoftContactsParams: A SoftContactsParams instance with the specified parameters.
-        """
-
-        return SoftContactsParams(
-            K=jnp.array(K, dtype=float),
-            D=jnp.array(D, dtype=float),
-            mu=jnp.array(mu, dtype=float),
-        )
-
-    @staticmethod
-    def build_default_from_physics_model(
-        physics_model: PhysicsModel,
-        static_friction_coefficient: jtp.FloatLike = 0.5,
-        max_penetration: jtp.FloatLike = 0.001,
-        number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
-        damping_ratio: jtp.FloatLike = 1.0,
-    ) -> SoftContactsParams:
-        """
-        Create a SoftContactsParams instance with good default parameters.
-
-        Args:
-            physics_model: The target physics model.
-            static_friction_coefficient: The static friction coefficient.
-            max_penetration: The maximum penetration depth.
-            number_of_active_collidable_points_steady_state: The number of contacts
-                supporting the weight of the model in steady state.
-            damping_ratio: The ratio controlling the damping behavior.
-
-        Returns:
-            A SoftContactsParams instance with the specified parameters.
-
-        Note:
-            The `damping_ratio` parameter allows to operate on the following conditions:
-            - ξ > 1.0: over-damped
-            - ξ = 1.0: critically damped
-            - ξ < 1.0: under-damped
-        """
-
-        # Use symbols for input parameters
-        ξ = damping_ratio
-        δ_max = max_penetration
-        μc = static_friction_coefficient
-
-        # Compute the total mass of the model
-        m = jnp.array(
-            [l.mass for l in physics_model.description.links_dict.values()]
-        ).sum()
-
-        # Extract gravity
-        g = -physics_model.gravity[0:3][-1]
-
-        # Compute the average support force on each collidable point
-        f_average = m * g / number_of_active_collidable_points_steady_state
-
-        # Compute the stiffness to get the desired steady-state penetration
-        K = f_average / jnp.power(δ_max, 3 / 2)
-
-        # Compute the damping using the damping ratio
-        critical_damping = 2 * jnp.sqrt(K * m)
-        D = ξ * critical_damping
-
-        return SoftContactsParams.build(K=K, D=D, mu=μc)
-
-
-@jax_dataclasses.pytree_dataclass
-class SoftContacts:
-    """Soft contacts model."""
-
-    parameters: SoftContactsParams = dataclasses.field(
-        default_factory=SoftContactsParams
-    )
-
-    terrain: Terrain = dataclasses.field(default_factory=FlatTerrain)
-
-    def contact_model(
-        self,
-        position: jtp.Vector,
-        velocity: jtp.Vector,
-        tangential_deformation: jtp.Vector,
-    ) -> Tuple[jtp.Vector, jtp.Vector]:
-        """
-        Compute the contact forces and material deformation rate.
-
-        Args:
-            position (jtp.Vector): The position of the collidable point.
-            velocity (jtp.Vector): The linear velocity of the collidable point.
-            tangential_deformation (jtp.Vector): The tangential deformation.
-
-        Returns:
-            Tuple[jtp.Vector, jtp.Vector]: A tuple containing the contact force and material deformation rate.
-        """
-
-        # Short name of parameters
-        K = self.parameters.K
-        D = self.parameters.D
-        μ = self.parameters.mu
-
-        # Material 3D tangential deformation and its derivative
-        m = tangential_deformation.squeeze()
-        ṁ = jnp.zeros_like(m)
-
-        # Note: all the small hardcoded tolerances in this method have been introduced
-        # to allow jax differentiating through this algorithm. They should not affect
-        # the accuracy of the simulation, although they might make it less readable.
-
-        # ========================
-        # Normal force computation
-        # ========================
-
-        # Unpack the position of the collidable point
-        px, py, pz = W_p_C = position.squeeze()
-        vx, vy, vz = W_ṗ_C = velocity.squeeze()
-
-        # Compute the terrain normal and the contact depth
-        n̂ = self.terrain.normal(x=px, y=py).squeeze()
-        h = jnp.array([0, 0, self.terrain.height(x=px, y=py) - pz])
-
-        # Compute the penetration depth normal to the terrain
-        δ = jnp.maximum(0.0, jnp.dot(h, n̂))
-
-        # Compute the penetration normal velocity
-        δ̇ = -jnp.dot(W_ṗ_C, n̂)
-
-        # Non-linear spring-damper model.
-        # This is the force magnitude along the direction normal to the terrain.
-        force_normal_mag = jax.lax.select(
-            pred=δ >= 1e-9,
-            on_true=jnp.sqrt(δ + 1e-12) * (K * δ + D * δ̇),
-            on_false=jnp.array(0.0),
-        )
-
-        # Prevent negative normal forces that might occur when δ̇ is largely negative
-        force_normal_mag = jnp.maximum(0.0, force_normal_mag)
-
-        # Compute the 3D linear force in C[W] frame
-        force_normal = force_normal_mag * n̂
-
-        # ====================================
-        # No friction and no tangential forces
-        # ====================================
-
-        # Compute the adjoint C[W]->W for transforming 6D forces from mixed to inertial.
-        # Note: this is equal to the 6D velocities transform: CW_X_W.transpose().
-        W_Xf_CW = jnp.vstack(
-            [
-                jnp.block([jnp.eye(3), jnp.zeros(shape=(3, 3))]),
-                jnp.block([Skew.wedge(W_p_C), jnp.eye(3)]),
-            ]
-        )
-
-        def with_no_friction():
-            # Compute 6D mixed force in C[W]
-            CW_f_lin = force_normal
-            CW_f = jnp.hstack([force_normal, jnp.zeros_like(CW_f_lin)])
-
-            # Compute lin-ang 6D forces (inertial representation)
-            W_f = W_Xf_CW @ CW_f
-
-            return W_f, ṁ
-
-        # =========================
-        # Compute tangential forces
-        # =========================
-
-        def with_friction():
-            # Initialize the tangential deformation rate ṁ.
-            # For inactive contacts with m≠0, this is the dynamics of the material
-            # relaxation converging exponentially to steady state.
-            ṁ = (-K / D) * m
-
-            # Check if the collidable point is below ground.
-            # Note: when δ=0, we consider the point still not it contact such that
-            #       we prevent divisions by 0 in the computations below.
-            active_contact = pz < self.terrain.height(x=px, y=py)
-
-            def above_terrain():
-                return jnp.zeros(6), ṁ
-
-            def below_terrain():
-                # Decompose the velocity in normal and tangential components
-                v_normal = jnp.dot(W_ṗ_C, n̂) * n̂
-                v_tangential = W_ṗ_C - v_normal
-
-                # Compute the tangential force. If inside the friction cone, the contact
-                f_tangential = -jnp.sqrt(δ + 1e-12) * (K * m + D * v_tangential)
-
-                def sticking_contact():
-                    # Sum the normal and tangential forces, and create the 6D force
-                    CW_f_stick = force_normal + f_tangential
-                    CW_f = jnp.hstack([CW_f_stick, jnp.zeros(3)])
-
-                    # In this case the 3D material deformation is the tangential velocity
-                    ṁ = v_tangential
-
-                    # Return the 6D force in the contact frame and
-                    # the deformation derivative
-                    return CW_f, ṁ
-
-                def slipping_contact():
-                    # Clip the tangential force if too small, allowing jax to
-                    # differentiate through the norm computation
-                    f_tangential_no_nan = jax.lax.select(
-                        pred=f_tangential.dot(f_tangential) >= 1e-9**2,
-                        on_true=f_tangential,
-                        on_false=jnp.array([1e-12, 0, 0]),
-                    )
-
-                    # Project the force to the friction cone boundary
-                    f_tangential_projected = (μ * force_normal_mag) * (
-                        f_tangential / jnp.linalg.norm(f_tangential_no_nan)
-                    )
-
-                    # Sum the normal and tangential forces, and create the 6D force
-                    CW_f_slip = force_normal + f_tangential_projected
-                    CW_f = jnp.hstack([CW_f_slip, jnp.zeros(3)])
-
-                    # Correct the material deformation derivative for slipping contacts.
-                    # Basically we compute ṁ such that we get `f_tangential` on the cone
-                    # given the current (m, δ).
-                    ε = 1e-9
-                    δε = jnp.maximum(δ, ε)
-                    α = -K * jnp.sqrt(δε)
-                    β = -D * jnp.sqrt(δε)
-                    ṁ = (f_tangential_projected - α * m) / β
-
-                    # Return the 6D force in the contact frame and
-                    # the deformation derivative
-                    return CW_f, ṁ
-
-                CW_f, ṁ = jax.lax.cond(
-                    pred=f_tangential.dot(f_tangential) > (μ * force_normal_mag) ** 2,
-                    true_fun=lambda _: slipping_contact(),
-                    false_fun=lambda _: sticking_contact(),
-                    operand=None,
-                )
-
-                # Express the 6D force in the world frame
-                W_f = W_Xf_CW @ CW_f
-
-                # Return the 6D force in the world frame and the deformation derivative
-                return W_f, ṁ
-
-            # (W_f, ṁ)
-            return jax.lax.cond(
-                pred=active_contact,
-                true_fun=lambda _: below_terrain(),
-                false_fun=lambda _: above_terrain(),
-                operand=None,
-            )
-
-        # (W_f, ṁ)
-        return jax.lax.cond(
-            pred=(μ == 0.0),
-            true_fun=lambda _: with_no_friction(),
-            false_fun=lambda _: with_friction(),
-            operand=None,
-        )

jaxsim/physics/algos/terrain.py

@@ -1,78 +0,0 @@
-import abc
-
-import jax.numpy as jnp
-import jax_dataclasses
-
-import jaxsim.typing as jtp
-
-
-class Terrain(abc.ABC):
-    delta = 0.010
-
-    @abc.abstractmethod
-    def height(self, x: float, y: float) -> float:
-        pass
-
-    def normal(self, x: float, y: float) -> jtp.Vector:
-        """
-        Compute the normal vector of the terrain at a specific (x, y) location.
-
-        Args:
-            x (float): The x-coordinate of the location.
-            y (float): The y-coordinate of the location.
-
-        Returns:
-            jtp.Vector: The normal vector of the terrain surface at the specified location.
-        """
-
-        # https://stackoverflow.com/a/5282364
-        h_xp = self.height(x=x + self.delta, y=y)
-        h_xm = self.height(x=x - self.delta, y=y)
-        h_yp = self.height(x=x, y=y + self.delta)
-        h_ym = self.height(x=x, y=y - self.delta)
-
-        n = jnp.array(
-            [(h_xm - h_xp) / (2 * self.delta), (h_ym - h_yp) / (2 * self.delta), 1.0]
-        )
-
-        return n / jnp.linalg.norm(n)
-
-
-@jax_dataclasses.pytree_dataclass
-class FlatTerrain(Terrain):
-    def height(self, x: float, y: float) -> float:
-        return 0.0
-
-
-@jax_dataclasses.pytree_dataclass
-class PlaneTerrain(Terrain):
-    plane_normal: list = jax_dataclasses.field(default_factory=lambda: [0, 0, 1.0])
-
-    @staticmethod
-    def build(plane_normal: list) -> "PlaneTerrain":
-        """
-        Create a PlaneTerrain instance with a specified plane normal vector.
-
-        Args:
-            plane_normal (list): The normal vector of the terrain plane.
-
-        Returns:
-            PlaneTerrain: A PlaneTerrain instance.
-        """
-
-        return PlaneTerrain(plane_normal=plane_normal)
-
-    def height(self, x: float, y: float) -> float:
-        """
-        Compute the height of the terrain at a specific (x, y) location on a plane.
-
-        Args:
-            x (float): The x-coordinate of the location.
-            y (float): The y-coordinate of the location.
-
-        Returns:
-            float: The height of the terrain at the specified location on the plane.
-        """
-
-        a, b, c = self.plane_normal
-        return -(a * x + b * y) / c

jaxsim/physics/algos/utils.py

@@ -1,69 +0,0 @@
-from typing import Tuple
-
-import jax.numpy as jnp
-
-import jaxsim.typing as jtp
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-
-def process_inputs(
-    physics_model: PhysicsModel,
-    xfb: jtp.Vector | None = None,
-    q: jtp.Vector | None = None,
-    qd: jtp.Vector | None = None,
-    qdd: jtp.Vector | None = None,
-    tau: jtp.Vector | None = None,
-    f_ext: jtp.Matrix | None = None,
-) -> Tuple[jtp.Vector, jtp.Vector, jtp.Vector, jtp.Vector, jtp.Vector, jtp.Matrix]:
-    """
-    Adjust the inputs to the physics model.
-
-    Args:
-        physics_model: The physics model.
-        xfb: The variables of the base link.
-        q: The generalized coordinates.
-        qd: The generalized velocities.
-        qdd: The generalized accelerations.
-        tau: The generalized forces.
-        f_ext: The link external forces.
-
-    Returns:
-        The adjusted inputs.
-    """
-
-    # Remove extra dimensions
-    q = q.squeeze() if q is not None else jnp.zeros(physics_model.dofs())
-    qd = qd.squeeze() if qd is not None else jnp.zeros(physics_model.dofs())
-    qdd = qdd.squeeze() if qdd is not None else jnp.zeros(physics_model.dofs())
-    tau = tau.squeeze() if tau is not None else jnp.zeros(physics_model.dofs())
-    xfb = xfb.squeeze() if xfb is not None else jnp.zeros(13).at[0].set(1)
-    f_ext = (
-        f_ext.squeeze()
-        if f_ext is not None
-        else jnp.zeros(shape=(physics_model.NB, 6)).squeeze()
-    )
-
-    # Fix case with just 1 DoF
-    q = jnp.atleast_1d(q)
-    qd = jnp.atleast_1d(qd)
-    qdd = jnp.atleast_1d(qdd)
-    tau = jnp.atleast_1d(tau)
-
-    # Fix case with just 1 body
-    f_ext = jnp.atleast_2d(f_ext)
-
-    # Validate dimensions
-    dofs = physics_model.dofs()
-
-    if xfb is not None and xfb.shape[0] != 13:
-        raise ValueError(xfb.shape)
-    if q is not None and q.shape[0] != dofs:
-        raise ValueError(q.shape, dofs)
-    if qd is not None and qd.shape[0] != dofs:
-        raise ValueError(qd.shape, dofs)
-    if tau is not None and tau.shape[0] != dofs:
-        raise ValueError(tau.shape, dofs)
-    if f_ext is not None and f_ext.shape != (physics_model.NB, 6):
-        raise ValueError(f_ext.shape, (physics_model.NB, 6))
-
-    return xfb, q, qd, qdd, tau, f_ext