jaxsim 0.1.dev401__py3-none-any.whl → 0.2.0__py3-none-any.whl

jaxsim/physics/algos/rnea_motors.py

@@ -1,196 +0,0 @@
-from typing import Tuple
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-import jaxsim.typing as jtp
-from jaxsim.math.adjoint import Adjoint
-from jaxsim.math.cross import Cross
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-from . import utils
-
-
-def rnea(
-    model: PhysicsModel,
-    xfb: jtp.Vector,
-    q: jtp.Vector,
-    qd: jtp.Vector,
-    qdd: jtp.Vector,
-    a0fb: jtp.Vector = jnp.zeros(6),
-    f_ext: jtp.Matrix | None = None,
-) -> Tuple[jtp.Vector, jtp.Vector]:
-    """
-    Recursive Newton-Euler Algorithm (RNEA) algorithm for inverse dynamics.
-    """
-
-    xfb, q, qd, qdd, _, f_ext = utils.process_inputs(
-        physics_model=model, xfb=xfb, q=q, qd=qd, qdd=qdd, f_ext=f_ext
-    )
-
-    a0fb = a0fb.squeeze()
-    gravity = model.gravity.squeeze()
-
-    if a0fb.shape[0] != 6:
-        raise ValueError(a0fb.shape)
-
-    M = model.spatial_inertias
-    pre_X_λi = model.tree_transforms
-    i_X_pre = model.joint_transforms(q=q)
-    S = model.motion_subspaces(q=q)
-    i_X_λi = jnp.zeros_like(pre_X_λi)
-
-    Γ = jnp.array([*model._joint_motor_gear_ratio.values()])
-    IM = jnp.array([*model._joint_motor_inertia.values()])
-    K_v = jnp.array([*model._joint_motor_viscous_friction.values()])
-    K̅ᵥ = jnp.diag(Γ.T * jnp.diag(K_v) * Γ)
-    m_S = jnp.concatenate([S[:1], S[1:] * Γ[:, None, None]], axis=0)
-
-    i_X_0 = jnp.zeros_like(pre_X_λi)
-    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
-
-    # Parent array mapping: i -> λ(i).
-    # Exception: λ(0) must not be used, it's initialized to -1.
-    λ = model.parent_array()
-
-    v = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    a = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    f = jnp.array([jnp.zeros([6, 1])] * model.NB)
-
-    v_m = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    a_m = jnp.array([jnp.zeros([6, 1])] * model.NB)
-    f_m = jnp.array([jnp.zeros([6, 1])] * model.NB)
-
-    # 6D transform of base velocity
-    B_X_W = Adjoint.from_quaternion_and_translation(
-        quaternion=xfb[0:4],
-        translation=xfb[4:7],
-        inverse=True,
-        normalize_quaternion=True,
-    )
-    i_X_λi = i_X_λi.at[0].set(B_X_W)
-
-    a_0 = -B_X_W @ jnp.vstack(gravity)
-    a = a.at[0].set(a_0)
-
-    if model.is_floating_base:
-        W_v_WB = jnp.vstack(jnp.hstack([xfb[10:13], xfb[7:10]]))
-
-        v_0 = B_X_W @ W_v_WB
-        v = v.at[0].set(v_0)
-
-        a_0 = B_X_W @ (jnp.vstack(a0fb) - jnp.vstack(gravity))
-        a = a.at[0].set(a_0)
-
-        f_0 = (
-            M[0] @ a[0]
-            + Cross.vx_star(v[0]) @ M[0] @ v[0]
-            - Adjoint.inverse(B_X_W).T @ jnp.vstack(f_ext[0])
-        )
-        f = f.at[0].set(f_0)
-
-    ForwardPassCarry = Tuple[
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-        jtp.MatrixJax,
-    ]
-    forward_pass_carry = (i_X_λi, v, v_m, a, a_m, i_X_0, f, f_m)
-
-    def forward_pass(
-        carry: ForwardPassCarry, i: jtp.Int
-    ) -> Tuple[ForwardPassCarry, None]:
-        ii = i - 1
-        i_X_λi, v, v_m, a, a_m, i_X_0, f, f_m = carry
-
-        vJ = S[i] * qd[ii]
-        vJ_m = m_S[i] * qd[ii]
-
-        i_X_λi_i = i_X_pre[i] @ pre_X_λi[i]
-        i_X_λi = i_X_λi.at[i].set(i_X_λi_i)
-
-        v_i = i_X_λi[i] @ v[λ[i]] + vJ
-        v = v.at[i].set(v_i)
-
-        v_i_m = i_X_λi[i] @ v_m[λ[i]] + vJ_m
-        v_m = v_m.at[i].set(v_i_m)
-
-        a_i = i_X_λi[i] @ a[λ[i]] + S[i] * qdd[ii] + Cross.vx(v[i]) @ vJ
-        a = a.at[i].set(a_i)
-
-        a_i_m = i_X_λi[i] @ a_m[λ[i]] + m_S[i] * qdd[ii] + Cross.vx(v_m[i]) @ vJ_m
-        a_m = a_m.at[i].set(a_i_m)
-
-        i_X_0_i = i_X_λi[i] @ i_X_0[λ[i]]
-        i_X_0 = i_X_0.at[i].set(i_X_0_i)
-        i_Xf_W = Adjoint.inverse(i_X_0[i] @ B_X_W).T
-
-        f_i = (
-            M[i] @ a[i]
-            + Cross.vx_star(v[i]) @ M[i] @ v[i]
-            - i_Xf_W @ jnp.vstack(f_ext[i])
-        )
-        f = f.at[i].set(f_i)
-
-        f_i_m = IM[i] * a_m[i] + Cross.vx_star(v_m[i]) * IM[i] @ v_m[i]
-        f_m = f_m.at[i].set(f_i_m)
-
-        return (i_X_λi, v, v_m, a, a_m, i_X_0, f, f_m), None
-
-    (i_X_λi, v, v_m, a, a_m, i_X_0, f, f_m), _ = jax.lax.scan(
-        f=forward_pass,
-        init=forward_pass_carry,
-        xs=np.arange(start=1, stop=model.NB),
-    )
-
-    tau = jnp.zeros_like(q)
-
-    BackwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax]
-    backward_pass_carry = (tau, f, f_m)
-
-    def backward_pass(
-        carry: BackwardPassCarry, i: jtp.Int
-    ) -> Tuple[BackwardPassCarry, None]:
-        ii = i - 1
-        tau, f, f_m = carry
-
-        value = S[i].T @ f[i] + m_S[i].T @ f_m[i]  # + K̅ᵥ[i] * qd[ii]
-        tau = tau.at[ii].set(value.squeeze())
-
-        def update_f(ffm: Tuple[jtp.MatrixJax, jtp.MatrixJax]) -> jtp.MatrixJax:
-            f, f_m = ffm
-            f_λi = f[λ[i]] + i_X_λi[i].T @ f[i]
-            f = f.at[λ[i]].set(f_λi)
-
-            f_m_λi = f_m[λ[i]] + i_X_λi[i].T @ f_m[i]
-            f_m = f_m.at[λ[i]].set(f_m_λi)
-            return f, f_m
-
-        f, f_m = jax.lax.cond(
-            pred=jnp.array([λ[i] != 0, model.is_floating_base]).any(),
-            true_fun=update_f,
-            false_fun=lambda f: f,
-            operand=(f, f_m),
-        )
-
-        return (tau, f, f_m), None
-
-    (tau, f, f_m), _ = jax.lax.scan(
-        f=backward_pass,
-        init=backward_pass_carry,
-        xs=np.flip(np.arange(start=1, stop=model.NB)),
-    )
-
-    # Handle 1 DoF models
-    tau = jnp.atleast_1d(tau.squeeze())
-    tau = jnp.vstack(tau) if tau.size > 0 else jnp.empty(shape=(0, 1))
-
-    # Express the base 6D force in the world frame
-    W_f0 = B_X_W.T @ jnp.vstack(f[0])
-
-    return W_f0, tau

jaxsim/physics/algos/soft_contacts.py

@@ -1,454 +0,0 @@
-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 . import utils
-
-
-@jax_dataclasses.pytree_dataclass
-class SoftContactsState:
-    """
-    State of the soft contacts model.
-
-    Attributes:
-        tangential_deformation (jtp.Matrix): The tangential deformation of the material at each collidable point.
-    """
-
-    tangential_deformation: jtp.Matrix
-
-    @staticmethod
-    def zero(
-        physics_model: jaxsim.physics.model.physics_model.PhysicsModel,
-    ) -> "SoftContactsState":
-        """
-        Modify the SoftContactsState instance imposing zero tangential deformation.
-
-        Args:
-            physics_model (jaxsim.physics.model.physics_model.PhysicsModel): The physics model.
-
-        Returns:
-            SoftContactsState: A SoftContactsState instance with zero tangential deformation.
-        """
-
-        return SoftContactsState(
-            tangential_deformation=jnp.zeros(shape=(3, physics_model.gc.body.size))
-        )
-
-    def valid(
-        self, physics_model: jaxsim.physics.model.physics_model.PhysicsModel
-    ) -> bool:
-        """
-        Check if the soft contacts state has valid shape.
-
-        Args:
-            physics_model (jaxsim.physics.model.physics_model.PhysicsModel): The physics model.
-
-        Returns:
-            bool: 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, physics_model.gc.body.size),
-            valid=True,
-        )
-
-    def replace(self, validate: bool = True, **kwargs) -> "SoftContactsState":
-        """
-        Replace attributes of the soft contacts state.
-
-        Args:
-            validate (bool, optional): Whether to validate the state after replacement. Defaults to True.
-
-        Returns:
-            SoftContactsState: A new SoftContactsState instance with replaced attributes.
-        """
-
-        with jax_dataclasses.copy_and_mutate(self, validate=validate) as updated_state:
-            _ = [updated_state.__setattr__(k, v) for k, v in kwargs.items()]
-
-        return updated_state
-
-
-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, model.gc.body
-    )
-
-    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: float = 1e6, D: float = 2_000, mu: float = 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),
-        )
-
-
-@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,
-        )