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

jaxsim/parsers/rod/utils.py

@@ -1,15 +1,24 @@
 import os
-from typing import Union
+import pathlib
+from collections.abc import Callable
+from typing import TypeVar
 
-import jax.numpy as jnp
 import numpy as np
 import numpy.typing as npt
 import rod
+import trimesh
+from rod.utils.resolve_uris import resolve_local_uri
 
+import jaxsim.typing as jtp
+from jaxsim import logging
+from jaxsim.math import Adjoint, Inertia
 from jaxsim.parsers import descriptions
+from jaxsim.parsers.rod import meshes
 
+MeshMappingMethod = TypeVar("MeshMappingMethod", bound=Callable[..., npt.NDArray])
 
-def from_sdf_inertial(inertial: rod.Inertial) -> npt.NDArray:
+
+def from_sdf_inertial(inertial: rod.Inertial) -> jtp.Matrix:
     """
     Extract the 6D inertia matrix from an SDF inertial element.
 
@@ -20,13 +29,10 @@ def from_sdf_inertial(inertial: rod.Inertial) -> npt.NDArray:
         The 6D inertia matrix of the link expressed in the link frame.
     """
 
-    from jaxsim.math.inertia import Inertia
-    from jaxsim.sixd import se3
-
-    # Extract the "mass" element
+    # Extract the "mass" element.
     m = inertial.mass
 
-    # Extract the "inertia" element
+    # Extract the "inertia" element.
     inertia_element = inertial.inertia
 
     ixx = inertia_element.ixx
@@ -36,7 +42,7 @@ def from_sdf_inertial(inertial: rod.Inertial) -> npt.NDArray:
     ixz = inertia_element.ixz if inertia_element.ixz is not None else 0.0
     iyz = inertia_element.iyz if inertia_element.iyz is not None else 0.0
 
-    # Build the 3x3 inertia matrix expressed in the CoM
+    # Build the 3x3 inertia matrix expressed in the CoM.
     I_CoM = np.array(
         [
             [ixx, ixy, ixz],
@@ -45,73 +51,52 @@ def from_sdf_inertial(inertial: rod.Inertial) -> npt.NDArray:
         ]
     )
 
-    # Build the 6x6 generalized inertia at the CoM
+    # Build the 6x6 generalized inertia at the CoM.
     M_CoM = Inertia.to_sixd(mass=m, com=np.zeros(3), I=I_CoM)
 
-    # Compute the transform from the inertial frame (CoM) to the link frame
+    # Compute the transform from the inertial frame (CoM) to the link frame.
     L_H_CoM = inertial.pose.transform() if inertial.pose is not None else np.eye(4)
 
-    # We need its inverse
-    CoM_H_L = se3.SE3.from_matrix(matrix=L_H_CoM).inverse()
-    CoM_X_L: npt.NDArray = CoM_H_L.adjoint()
+    # We need its inverse.
+    CoM_X_L = Adjoint.from_transform(transform=L_H_CoM, inverse=True)
 
-    # Express the CoM inertia matrix in the link frame L
+    # Express the CoM inertia matrix in the link frame L.
     M_L = CoM_X_L.T @ M_CoM @ CoM_X_L
 
-    return jnp.array(M_L)
+    return M_L.astype(dtype=float)
 
 
-def axis_to_jtype(
-    axis: rod.Axis, type: str
-) -> Union[descriptions.JointType, descriptions.JointDescriptor]:
+def joint_to_joint_type(joint: rod.Joint) -> int:
     """
-    Convert an SDF axis to a joint type.
+    Extract the joint type from an SDF joint.
 
     Args:
-        axis: The SDF axis.
-        type: The SDF joint type.
+        joint: The parsed SDF joint.
 
     Returns:
-        The corresponding joint type description.
+        The integer corresponding to the joint type.
     """
 
-    if type == "fixed":
-        return descriptions.JointType.F
+    axis = joint.axis
+    joint_type = joint.type
+
+    if joint_type == "fixed":
+        return descriptions.JointType.Fixed
 
     if not (axis.xyz is not None and axis.xyz.xyz is not None):
         raise ValueError("Failed to read axis xyz data")
 
-    axis_xyz = np.array(axis.xyz.xyz)
-
-    if np.allclose(axis_xyz, [1, 0, 0]) and type in {"revolute", "continuous"}:
-        return descriptions.JointType.Rx
-
-    if np.allclose(axis_xyz, [0, 1, 0]) and type in {"revolute", "continuous"}:
-        return descriptions.JointType.Ry
-
-    if np.allclose(axis_xyz, [0, 0, 1]) and type in {"revolute", "continuous"}:
-        return descriptions.JointType.Rz
+    # Make sure that the axis is a unary vector.
+    axis_xyz = np.array(axis.xyz.xyz).astype(float)
+    axis_xyz = axis_xyz / np.linalg.norm(axis_xyz)
 
-    if np.allclose(axis_xyz, [1, 0, 0]) and type == "prismatic":
-        return descriptions.JointType.Px
+    if joint_type in {"revolute", "continuous"}:
+        return descriptions.JointType.Revolute
 
-    if np.allclose(axis_xyz, [0, 1, 0]) and type == "prismatic":
-        return descriptions.JointType.Py
+    if joint_type == "prismatic":
+        return descriptions.JointType.Prismatic
 
-    if np.allclose(axis_xyz, [0, 0, 1]) and type == "prismatic":
-        return descriptions.JointType.Pz
-
-    if type == "revolute":
-        return descriptions.JointGenericAxis(
-            code=descriptions.JointType.R, axis=np.array(axis_xyz, dtype=float)
-        )
-
-    if type == "prismatic":
-        return descriptions.JointGenericAxis(
-            code=descriptions.JointType.P, axis=np.array(axis_xyz, dtype=float)
-        )
-
-    raise ValueError("Joint not supported", axis_xyz, type)
+    raise ValueError("Joint not supported", axis_xyz, joint_type)
 
 
 def create_box_collision(
@@ -132,22 +117,19 @@ def create_box_collision(
 
     center = np.array([x / 2, y / 2, z / 2])
 
-    box_corners = (
-        np.vstack(
-            [
-                np.array([0, 0, 0]),
-                np.array([x, 0, 0]),
-                np.array([x, y, 0]),
-                np.array([0, y, 0]),
-                np.array([0, 0, z]),
-                np.array([x, 0, z]),
-                np.array([x, y, z]),
-                np.array([0, y, z]),
-            ]
-        )
-        - center
+    # Define the bottom corners.
+    bottom_corners = np.array([[0, 0, 0], [x, 0, 0], [x, y, 0], [0, y, 0]])
+
+    # Conditionally add the top corners based on the environment variable.
+    top_corners = (
+        np.array([[0, 0, z], [x, 0, z], [x, y, z], [0, y, z]])
+        if not os.environ.get("JAXSIM_COLLISION_USE_BOTTOM_ONLY", "0")
+        else []
     )
 
+    # Combine and shift by the center
+    box_corners = np.vstack([bottom_corners, *top_corners]) - center
+
     H = collision.pose.transform() if collision.pose is not None else np.eye(4)
 
     center_wrt_link = (H @ np.hstack([center, 1.0]))[0:-1]
@@ -158,7 +140,7 @@ def create_box_collision(
     collidable_points = [
         descriptions.CollidablePoint(
             parent_link=link_description,
-            position=corner,
+            position=np.array(corner),
             enabled=True,
         )
         for corner in box_corners_wrt_link.T
@@ -185,25 +167,33 @@ def create_sphere_collision(
 
     # From https://stackoverflow.com/a/26127012
     def fibonacci_sphere(samples: int) -> npt.NDArray:
-        points = []
-        phi = np.pi * (3.0 - np.sqrt(5.0))  # golden angle in radians
-
-        for i in range(samples):
-            y = 1 - (i / float(samples - 1)) * 2  # y goes from 1 to -1
-            radius = np.sqrt(1 - y * y)  # radius at y
-
-            theta = phi * i  # golden angle increment
-
-            x = np.cos(theta) * radius
-            z = np.sin(theta) * radius
+        # Get the golden ratio in radians.
+        phi = np.pi * (3.0 - np.sqrt(5.0))
+
+        # Generate the points.
+        points = [
+            np.array(
+                [
+                    np.cos(phi * i)
+                    * np.sqrt(1 - (y := 1 - 2 * i / (samples - 1)) ** 2),
+                    y,
+                    np.sin(phi * i) * np.sqrt(1 - y**2),
+                ]
+            )
+            for i in range(samples)
+        ]
 
-            points.append(np.array([x, y, z]))
+        # Filter to keep only the bottom half if required.
+        if os.environ.get("JAXSIM_COLLISION_USE_BOTTOM_ONLY", "0"):
+            # Keep only the points with z <= 0.
+            points = [point for point in points if point[2] <= 0]
 
         return np.vstack(points)
 
     r = collision.geometry.sphere.radius
+
     sphere_points = r * fibonacci_sphere(
-        samples=int(os.getenv(key="JAXSIM_COLLISION_SPHERE_POINTS", default="250"))
+        samples=int(os.getenv(key="JAXSIM_COLLISION_SPHERE_POINTS", default="50"))
     )
 
     H = collision.pose.transform() if collision.pose is not None else np.eye(4)
@@ -217,7 +207,7 @@ def create_sphere_collision(
     collidable_points = [
         descriptions.CollidablePoint(
             parent_link=link_description,
-            position=point,
+            position=np.array(point),
             enabled=True,
         )
         for point in sphere_points_wrt_link.T
@@ -226,3 +216,58 @@ def create_sphere_collision(
     return descriptions.SphereCollision(
         collidable_points=collidable_points, center=center_wrt_link
     )
+
+
+def create_mesh_collision(
+    collision: rod.Collision,
+    link_description: descriptions.LinkDescription,
+    method: MeshMappingMethod = None,
+) -> descriptions.MeshCollision:
+    """
+    Create a mesh collision from an SDF collision element.
+
+    Args:
+        collision: The SDF collision element.
+        link_description: The link description.
+        method: The method to use for mesh wrapping.
+
+    Returns:
+        The mesh collision description.
+    """
+
+    file = pathlib.Path(resolve_local_uri(uri=collision.geometry.mesh.uri))
+    file_type = file.suffix.replace(".", "")
+    mesh = trimesh.load_mesh(file, file_type=file_type)
+
+    if mesh.is_empty:
+        raise RuntimeError(f"Failed to process '{file}' with trimesh")
+
+    mesh.apply_scale(collision.geometry.mesh.scale)
+    logging.info(
+        msg=f"Loading mesh {collision.geometry.mesh.uri} with scale {collision.geometry.mesh.scale}, file type '{file_type}'"
+    )
+
+    if method is None:
+        method = meshes.VertexExtraction()
+        logging.debug("Using default Vertex Extraction method for mesh wrapping")
+    else:
+        logging.debug(f"Using method {method} for mesh wrapping")
+
+    points = method(mesh=mesh)
+    logging.debug(f"Extracted {len(points)} points from mesh")
+
+    W_H_L = collision.pose.transform() if collision.pose is not None else np.eye(4)
+
+    # Extract translation from transformation matrix
+    W_p_L = W_H_L[:3, 3]
+    mesh_points_wrt_link = points @ W_H_L[:3, :3].T + W_p_L
+    collidable_points = [
+        descriptions.CollidablePoint(
+            parent_link=link_description,
+            position=point,
+            enabled=True,
+        )
+        for point in mesh_points_wrt_link
+    ]
+
+    return descriptions.MeshCollision(collidable_points=collidable_points, center=W_p_L)

jaxsim/rbda/__init__.py

@@ -0,0 +1,11 @@
+from . import contacts
+from .aba import aba
+from .collidable_points import collidable_points_pos_vel
+from .crba import crba
+from .forward_kinematics import forward_kinematics, forward_kinematics_model
+from .jacobian import (
+    jacobian,
+    jacobian_derivative_full_doubly_left,
+    jacobian_full_doubly_left,
+)
+from .rnea import rnea

jaxsim/rbda/aba.py

@@ -0,0 +1,289 @@
+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, StandardGravity
+
+from . import utils
+
+
+def aba(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.VectorLike,
+    base_quaternion: jtp.VectorLike,
+    joint_positions: jtp.VectorLike,
+    base_linear_velocity: jtp.VectorLike,
+    base_angular_velocity: jtp.VectorLike,
+    joint_velocities: jtp.VectorLike,
+    joint_forces: jtp.VectorLike | None = None,
+    link_forces: jtp.MatrixLike | None = None,
+    standard_gravity: jtp.FloatLike = StandardGravity,
+) -> tuple[jtp.Vector, jtp.Vector]:
+    """
+    Compute forward dynamics using the Articulated Body Algorithm (ABA).
+
+    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.
+        joint_forces: The forces applied to 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 base acceleration in inertial-fixed representation
+        and the joint accelerations that result from the applications of the given
+        joint and link forces.
+
+    Note:
+        The algorithm expects a quaternion with unit norm.
+    """
+
+    W_p_B, W_Q_B, s, W_v_WB, ṡ, _, _, τ, 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=None,
+        base_angular_acceleration=None,
+        joint_accelerations=None,
+        joint_forces=joint_forces,
+        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
+
+    # 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(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))
+    c = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    pA = jnp.zeros(shape=(model.number_of_links(), 6, 1))
+    MA = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+
+    # 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 base quantities.
+    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)
+
+        # Initialize the articulated-body inertia (Mᴬ) of base link.
+        MA_0 = M[0]
+        MA = MA.at[0].set(MA_0)
+
+        # Initialize the articulated-body bias force (pᴬ) of the base link.
+        pA_0 = Cross.vx_star(v[0]) @ MA[0] @ v[0] - W_X_B.T @ jnp.vstack(W_f[0])
+        pA = pA.at[0].set(pA_0)
+
+    # ======
+    # Pass 1
+    # ======
+
+    Pass1Carry = tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix]
+    pass_1_carry: Pass1Carry = (v, c, MA, pA, i_X_0)
+
+    # Propagate kinematics and initialize AB inertia and AB bias forces.
+    def loop_body_pass1(carry: Pass1Carry, i: jtp.Int) -> tuple[Pass1Carry, None]:
+
+        ii = i - 1
+        v, c, MA, pA, i_X_0 = 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)
+
+        c_i = Cross.vx(v[i]) @ vJ
+        c = c.at[i].set(c_i)
+
+        # Initialize the articulated-body inertia.
+        MA_i = jnp.array(M[i])
+        MA = MA.at[i].set(MA_i)
+
+        # Compute the link-to-base transform.
+        i_Xi_0 = i_X_λi[i] @ i_X_0[λ[i]]
+        i_X_0 = i_X_0.at[i].set(i_Xi_0)
+
+        # Compute link-to-world transform for the 6D force.
+        i_Xf_W = Adjoint.inverse(i_X_0[i] @ B_X_W).T
+
+        # Initialize articulated-body bias force.
+        pA_i = Cross.vx_star(v[i]) @ M[i] @ v[i] - i_Xf_W @ jnp.vstack(W_f[i])
+        pA = pA.at[i].set(pA_i)
+
+        return (v, c, MA, pA, i_X_0), None
+
+    (v, c, MA, pA, i_X_0), _ = (
+        jax.lax.scan(
+            f=loop_body_pass1,
+            init=pass_1_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(v, c, MA, pA, i_X_0), None]
+    )
+
+    # ======
+    # Pass 2
+    # ======
+
+    U = jnp.zeros_like(S)
+    d = jnp.zeros(shape=(model.number_of_links(), 1))
+    u = jnp.zeros(shape=(model.number_of_links(), 1))
+
+    Pass2Carry = tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix]
+    pass_2_carry: Pass2Carry = (U, d, u, MA, pA)
+
+    def loop_body_pass2(carry: Pass2Carry, i: jtp.Int) -> tuple[Pass2Carry, None]:
+
+        ii = i - 1
+        U, d, u, MA, pA = carry
+
+        U_i = MA[i] @ S[i]
+        U = U.at[i].set(U_i)
+
+        d_i = S[i].T @ U[i]
+        d = d.at[i].set(d_i.squeeze())
+
+        u_i = τ[ii] - S[i].T @ pA[i]
+        u = u.at[i].set(u_i.squeeze())
+
+        # Compute the articulated-body inertia and bias force of this link.
+        Ma = MA[i] - U[i] / d[i] @ U[i].T
+        pa = pA[i] + Ma @ c[i] + U[i] * (u[i] / d[i])
+
+        # Propagate them to the parent, handling the base link.
+        def propagate(
+            MA_pA: tuple[jtp.Matrix, jtp.Matrix]
+        ) -> tuple[jtp.Matrix, jtp.Matrix]:
+
+            MA, pA = MA_pA
+
+            MA_λi = MA[λ[i]] + i_X_λi[i].T @ Ma @ i_X_λi[i]
+            MA = MA.at[λ[i]].set(MA_λi)
+
+            pA_λi = pA[λ[i]] + i_X_λi[i].T @ pa
+            pA = pA.at[λ[i]].set(pA_λi)
+
+            return MA, pA
+
+        MA, pA = jax.lax.cond(
+            pred=jnp.logical_or(λ[i] != 0, model.floating_base()),
+            true_fun=propagate,
+            false_fun=lambda MA_pA: MA_pA,
+            operand=(MA, pA),
+        )
+
+        return (U, d, u, MA, pA), None
+
+    (U, d, u, MA, pA), _ = (
+        jax.lax.scan(
+            f=loop_body_pass2,
+            init=pass_2_carry,
+            xs=jnp.flip(jnp.arange(start=1, stop=model.number_of_links())),
+        )
+        if model.number_of_links() > 1
+        else [(U, d, u, MA, pA), None]
+    )
+
+    # ======
+    # Pass 3
+    # ======
+
+    if model.floating_base():
+        a0 = jnp.linalg.solve(-MA[0], pA[0])
+    else:
+        a0 = -B_X_W @ W_g
+
+    s̈ = jnp.zeros_like(s)
+    a = jnp.zeros_like(v).at[0].set(a0)
+
+    Pass3Carry = tuple[jtp.Matrix, jtp.Vector]
+    pass_3_carry = (a, s̈)
+
+    def loop_body_pass3(carry: Pass3Carry, i: jtp.Int) -> tuple[Pass3Carry, None]:
+
+        ii = i - 1
+        a, s̈ = carry
+
+        # Propagate the link acceleration.
+        a_i = i_X_λi[i] @ a[λ[i]] + c[i]
+
+        # Compute the joint acceleration.
+        s̈_ii = (u[i] - U[i].T @ a_i) / d[i]
+        s̈ = s̈.at[ii].set(s̈_ii.squeeze())
+
+        # Sum the joint acceleration to the parent link acceleration.
+        a_i = a_i + S[i] * s̈[ii]
+        a = a.at[i].set(a_i)
+
+        return (a, s̈), None
+
+    (a, s̈), _ = (
+        jax.lax.scan(
+            f=loop_body_pass3,
+            init=pass_3_carry,
+            xs=jnp.arange(1, model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(a, s̈), None]
+    )
+
+    # ==============
+    # Adjust outputs
+    # ==============
+
+    # TODO: remove vstack and shape=(6, 1)?
+    if model.floating_base():
+        # Convert the base acceleration to inertial-fixed representation,
+        # and add gravity.
+        B_a_WB = a[0]
+        W_a_WB = W_X_B @ B_a_WB + W_g
+    else:
+        W_a_WB = jnp.zeros(6)
+
+    return W_a_WB.squeeze(), jnp.atleast_1d(s̈.squeeze())