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

jaxsim/rbda/rnea.py

@@ -0,0 +1,237 @@
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Cross, Quaternion, StandardGravity
+
+from . import utils
+
+
+def rnea(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.Vector,
+    base_quaternion: jtp.Vector,
+    joint_positions: jtp.Vector,
+    base_linear_velocity: jtp.Vector,
+    base_angular_velocity: jtp.Vector,
+    joint_velocities: jtp.Vector,
+    base_linear_acceleration: jtp.Vector | None = None,
+    base_angular_acceleration: jtp.Vector | None = None,
+    joint_accelerations: jtp.Vector | None = None,
+    link_forces: jtp.Matrix | None = None,
+    standard_gravity: jtp.FloatLike = StandardGravity,
+) -> Tuple[jtp.Vector, jtp.Vector]:
+    """
+    Compute inverse dynamics using the Recursive Newton-Euler Algorithm (RNEA).
+
+    Args:
+        model: The model to consider.
+        base_position: The position of the base link.
+        base_quaternion: The quaternion of the base link.
+        joint_positions: The positions of the joints.
+        base_linear_velocity:
+            The linear velocity of the base link in inertial-fixed representation.
+        base_angular_velocity:
+            The angular velocity of the base link in inertial-fixed representation.
+        joint_velocities: The velocities of the joints.
+        base_linear_acceleration:
+            The linear acceleration of the base link in inertial-fixed representation.
+        base_angular_acceleration:
+            The angular acceleration of the base link in inertial-fixed representation.
+        joint_accelerations: The accelerations of the joints.
+        link_forces:
+            The forces applied to the links expressed in the world frame.
+        standard_gravity: The standard gravity constant.
+
+    Returns:
+        A tuple containing the 6D force applied to the base link expressed in the
+        world frame and the joint forces that, when applied respectively to the base
+        link and joints, produce the given base and joint accelerations.
+    """
+
+    W_p_B, W_Q_B, s, W_v_WB, ṡ, W_v̇_WB, s̈, _, W_f, W_g = utils.process_inputs(
+        model=model,
+        base_position=base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=joint_positions,
+        base_linear_velocity=base_linear_velocity,
+        base_angular_velocity=base_angular_velocity,
+        joint_velocities=joint_velocities,
+        base_linear_acceleration=base_linear_acceleration,
+        base_angular_acceleration=base_angular_acceleration,
+        joint_accelerations=joint_accelerations,
+        link_forces=link_forces,
+        standard_gravity=standard_gravity,
+    )
+
+    W_g = jnp.atleast_2d(W_g).T
+    W_v_WB = jnp.atleast_2d(W_v_WB).T
+    W_v̇_WB = jnp.atleast_2d(W_v̇_WB).T
+
+    # Get the 6D spatial inertia matrices of all links.
+    M = js.model.link_spatial_inertia_matrices(model=model)
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute the base transform.
+    W_H_B = jaxlie.SE3.from_rotation_and_translation(
+        rotation=jaxlie.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(wxyz=W_Q_B)),
+        translation=W_p_B,
+    )
+
+    # Compute 6D transforms of the base velocity.
+    W_X_B = W_H_B.adjoint()
+    B_X_W = W_H_B.inverse().adjoint()
+
+    # Compute the parent-to-child adjoints and the motion subspaces of the joints.
+    # These transforms define the relative kinematics of the entire model, including
+    # the base transform for both floating-base and fixed-base models.
+    i_X_λi, S = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
+        joint_positions=s, base_transform=W_H_B.as_matrix()
+    )
+
+    # Allocate buffers.
+    v = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    a = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    f = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+
+    # Allocate the buffer of transforms link -> base.
+    i_X_0 = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    i_X_0 = i_X_0.at[0].set(jnp.eye(6))
+
+    # Initialize the acceleration of the base link.
+    a_0 = -B_X_W @ W_g
+    a = a.at[0].set(a_0)
+
+    if model.floating_base():
+
+        # Base velocity v₀ in body-fixed representation.
+        v_0 = B_X_W @ W_v_WB
+        v = v.at[0].set(v_0)
+
+        # Base acceleration a₀ in body-fixed representation w/o gravity.
+        a_0 = B_X_W @ (W_v̇_WB - W_g)
+        a = a.at[0].set(a_0)
+
+        # Force applied to the base link that produce the base acceleration w/o gravity.
+        f_0 = (
+            M[0] @ a[0]
+            + Cross.vx_star(v[0]) @ M[0] @ v[0]
+            - W_X_B.T @ jnp.vstack(W_f[0])
+        )
+        f = f.at[0].set(f_0)
+
+    # ======
+    # Pass 1
+    # ======
+
+    ForwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax]
+    forward_pass_carry: ForwardPassCarry = (v, a, i_X_0, f)
+
+    def forward_pass(
+        carry: ForwardPassCarry, i: jtp.Int
+    ) -> Tuple[ForwardPassCarry, None]:
+
+        ii = i - 1
+        v, a, i_X_0, f = carry
+
+        # Project the joint velocity into its motion subspace.
+        vJ = S[i] * ṡ[ii]
+
+        # Propagate the link velocity.
+        v_i = i_X_λi[i] @ v[λ[i]] + vJ
+        v = v.at[i].set(v_i)
+
+        # Propagate the link acceleration.
+        a_i = i_X_λi[i] @ a[λ[i]] + S[i] * s̈[ii] + Cross.vx(v[i]) @ vJ
+        a = a.at[i].set(a_i)
+
+        # Compute the link-to-base transform.
+        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)
+
+        # Compute link-to-world transform for the 6D force.
+        i_Xf_W = Adjoint.inverse(i_X_0[i] @ B_X_W).T
+
+        # Compute the force acting on the link.
+        f_i = (
+            M[i] @ a[i]
+            + Cross.vx_star(v[i]) @ M[i] @ v[i]
+            - i_Xf_W @ jnp.vstack(W_f[i])
+        )
+        f = f.at[i].set(f_i)
+
+        return (v, a, i_X_0, f), None
+
+    (v, a, i_X_0, f), _ = (
+        jax.lax.scan(
+            f=forward_pass,
+            init=forward_pass_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(v, a, i_X_0, f), None]
+    )
+
+    # ======
+    # Pass 2
+    # ======
+
+    τ = jnp.zeros_like(s)
+
+    BackwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
+    backward_pass_carry: BackwardPassCarry = (τ, f)
+
+    def backward_pass(
+        carry: BackwardPassCarry, i: jtp.Int
+    ) -> Tuple[BackwardPassCarry, None]:
+
+        ii = i - 1
+        τ, f = carry
+
+        # Project the 6D force to the DoF of the joint.
+        τ_i = S[i].T @ f[i]
+        τ = τ.at[ii].set(τ_i.squeeze())
+
+        # Propagate the force to the parent link.
+        def update_f(f: jtp.MatrixJax) -> jtp.MatrixJax:
+
+            f_λi = f[λ[i]] + i_X_λi[i].T @ f[i]
+            f = f.at[λ[i]].set(f_λi)
+
+            return f
+
+        f = jax.lax.cond(
+            pred=jnp.logical_or(λ[i] != 0, model.floating_base()),
+            true_fun=update_f,
+            false_fun=lambda f: f,
+            operand=f,
+        )
+
+        return (τ, f), None
+
+    (τ, f), _ = (
+        jax.lax.scan(
+            f=backward_pass,
+            init=backward_pass_carry,
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        )
+        if model.number_of_links() > 1
+        else [(τ, f), None]
+    )
+
+    # ==============
+    # Adjust outputs
+    # ==============
+
+    # Express the base 6D force in the world frame.
+    W_f0 = B_X_W.T @ f[0]
+
+    return W_f0.squeeze(), jnp.atleast_1d(τ.squeeze())

jaxsim/rbda/soft_contacts.py

@@ -0,0 +1,296 @@
+from __future__ import annotations
+
+import dataclasses
+
+import jax
+import jax.numpy as jnp
+import jax_dataclasses
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import Skew, StandardGravity
+from jaxsim.terrain import FlatTerrain, Terrain
+from jaxsim.utils import JaxsimDataclass
+
+
+@jax_dataclasses.pytree_dataclass
+class SoftContactsParams(JaxsimDataclass):
+    """Parameters of the soft contacts model."""
+
+    K: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(1e6, dtype=float)
+    )
+
+    D: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array(2000, dtype=float)
+    )
+
+    mu: jtp.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: The stiffness parameter.
+            D: The damping parameter of the soft contacts model.
+            mu: The static friction coefficient.
+
+        Returns:
+            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_jaxsim_model(
+        model: js.model.JaxSimModel,
+        *,
+        standard_gravity: jtp.FloatLike = StandardGravity,
+        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:
+            model: The target model.
+            standard_gravity: The standard gravity constant.
+            static_friction_coefficient:
+                The static friction coefficient between the model and the terrain.
+            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(model.kin_dyn_parameters.link_parameters.mass).sum()
+
+        # Rename the standard gravity
+        g = standard_gravity
+
+        # 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: The position of the collidable point.
+            velocity: The linear velocity of the collidable point.
+            tangential_deformation: The tangential deformation.
+
+        Returns:
+            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():
+                    # Project the force to the friction cone boundary
+                    f_tangential_projected = (μ * force_normal_mag) * (
+                        f_tangential / jnp.maximum(jnp.linalg.norm(f_tangential), 1e-9)
+                    )
+
+                    # 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/rbda/utils.py

@@ -0,0 +1,152 @@
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import StandardGravity
+
+
+def process_inputs(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.VectorLike | None = None,
+    base_quaternion: jtp.VectorLike | None = None,
+    joint_positions: jtp.VectorLike | None = None,
+    base_linear_velocity: jtp.VectorLike | None = None,
+    base_angular_velocity: jtp.VectorLike | None = None,
+    joint_velocities: jtp.VectorLike | None = None,
+    base_linear_acceleration: jtp.VectorLike | None = None,
+    base_angular_acceleration: jtp.VectorLike | None = None,
+    joint_accelerations: jtp.VectorLike | None = None,
+    joint_forces: jtp.VectorLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
+    standard_gravity: jtp.VectorLike | None = None,
+) -> tuple[
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Vector,
+    jtp.Matrix,
+    jtp.Vector,
+]:
+    """
+    Adjust the inputs to rigid-body dynamics algorithms.
+
+    Args:
+        model: The model to consider.
+        base_position: The position of the base link.
+        base_quaternion: The quaternion of the base link.
+        joint_positions: The positions of the joints.
+        base_linear_velocity: The linear velocity of the base link.
+        base_angular_velocity: The angular velocity of the base link.
+        joint_velocities: The velocities of the joints.
+        base_linear_acceleration: The linear acceleration of the base link.
+        base_angular_acceleration: The angular acceleration of the base link.
+        joint_accelerations: The accelerations of the joints.
+        joint_forces: The forces applied to the joints.
+        link_forces: The forces applied to the links.
+        standard_gravity: The standard gravity constant.
+
+    Returns:
+        The adjusted inputs.
+    """
+
+    dofs = model.dofs()
+    nl = model.number_of_links()
+
+    # Floating-base position.
+    W_p_B = base_position
+    W_Q_B = base_quaternion
+    s = joint_positions
+
+    # Floating-base velocity in inertial-fixed representation.
+    W_vl_WB = base_linear_velocity
+    W_ω_WB = base_angular_velocity
+    ṡ = joint_velocities
+
+    # Floating-base acceleration in inertial-fixed representation.
+    W_v̇l_WB = base_linear_acceleration
+    W_ω̇_WB = base_angular_acceleration
+    s̈ = joint_accelerations
+
+    # System dynamics inputs.
+    f = link_forces
+    τ = joint_forces
+
+    # Fill missing data and adjust dimensions.
+    s = jnp.atleast_1d(s.squeeze()) if s is not None else jnp.zeros(dofs)
+    ṡ = jnp.atleast_1d(ṡ.squeeze()) if ṡ is not None else jnp.zeros(dofs)
+    s̈ = jnp.atleast_1d(s̈.squeeze()) if s̈ is not None else jnp.zeros(dofs)
+    τ = jnp.atleast_1d(τ.squeeze()) if τ is not None else jnp.zeros(dofs)
+    W_vl_WB = jnp.atleast_1d(W_vl_WB.squeeze()) if W_vl_WB is not None else jnp.zeros(3)
+    W_v̇l_WB = jnp.atleast_1d(W_v̇l_WB.squeeze()) if W_v̇l_WB is not None else jnp.zeros(3)
+    W_p_B = jnp.atleast_1d(W_p_B.squeeze()) if W_p_B is not None else jnp.zeros(3)
+    W_ω_WB = jnp.atleast_1d(W_ω_WB.squeeze()) if W_ω_WB is not None else jnp.zeros(3)
+    W_ω̇_WB = jnp.atleast_1d(W_ω̇_WB.squeeze()) if W_ω̇_WB is not None else jnp.zeros(3)
+    f = jnp.atleast_2d(f.squeeze()) if f is not None else jnp.zeros(shape=(nl, 6))
+    W_Q_B = (
+        jnp.atleast_1d(W_Q_B.squeeze())
+        if W_Q_B is not None
+        else jnp.array([1.0, 0, 0, 0])
+    )
+    standard_gravity = (
+        jnp.array(standard_gravity).squeeze()
+        if standard_gravity is not None
+        else StandardGravity
+    )
+
+    if s.shape != (dofs,):
+        raise ValueError(s.shape, dofs)
+
+    if ṡ.shape != (dofs,):
+        raise ValueError(ṡ.shape, dofs)
+
+    if s̈.shape != (dofs,):
+        raise ValueError(s̈.shape, dofs)
+
+    if τ.shape != (dofs,):
+        raise ValueError(τ.shape, dofs)
+
+    if W_p_B.shape != (3,):
+        raise ValueError(W_p_B.shape, (3,))
+
+    if W_vl_WB.shape != (3,):
+        raise ValueError(W_vl_WB.shape, (3,))
+
+    if W_ω_WB.shape != (3,):
+        raise ValueError(W_ω_WB.shape, (3,))
+
+    if W_v̇l_WB.shape != (3,):
+        raise ValueError(W_v̇l_WB.shape, (3,))
+
+    if W_ω̇_WB.shape != (3,):
+        raise ValueError(W_ω̇_WB.shape, (3,))
+
+    if f.shape != (nl, 6):
+        raise ValueError(f.shape, (nl, 6))
+
+    if W_Q_B.shape != (4,):
+        raise ValueError(W_Q_B.shape, (4,))
+
+    # Pack the 6D base velocity and acceleration.
+    W_v_WB = jnp.hstack([W_vl_WB, W_ω_WB])
+    W_v̇_WB = jnp.hstack([W_v̇l_WB, W_ω̇_WB])
+
+    # Create the 6D gravity acceleration.
+    W_g = jnp.zeros(6).at[2].set(-standard_gravity)
+
+    return (
+        W_p_B.astype(float),
+        W_Q_B.astype(float),
+        s.astype(float),
+        W_v_WB.astype(float),
+        ṡ.astype(float),
+        W_v̇_WB.astype(float),
+        s̈.astype(float),
+        τ.astype(float),
+        f.astype(float),
+        W_g.astype(float),
+    )

jaxsim/terrain/__init__.py

@@ -0,0 +1,2 @@
+from . import terrain
+from .terrain import FlatTerrain, PlaneTerrain, Terrain