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

jaxsim/parsers/rod/parser.py

@@ -10,7 +10,7 @@ from jaxsim import logging
 from jaxsim.math.quaternion import Quaternion
 from jaxsim.parsers import descriptions, kinematic_graph
 
-from . import utils as utils
+from . import utils
 
 
 class SDFData(NamedTuple):
@@ -135,11 +135,13 @@ def extract_model_data(
                 parent=world_link,
                 child=links_dict[j.child],
                 jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
-                axis=np.array(j.axis.xyz.xyz)
-                if j.axis is not None
-                and j.axis.xyz is not None
-                and j.axis.xyz.xyz is not None
-                else None,
+                axis=(
+                    np.array(j.axis.xyz.xyz)
+                    if j.axis is not None
+                    and j.axis.xyz is not None
+                    and j.axis.xyz.xyz is not None
+                    else None
+                ),
                 pose=j.pose.transform() if j.pose is not None else np.eye(4),
             )
             for j in sdf_model.joints()
@@ -200,41 +202,55 @@ def extract_model_data(
             parent=links_dict[j.parent],
             child=links_dict[j.child],
             jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
-            axis=np.array(j.axis.xyz.xyz)
-            if j.axis is not None
-            and j.axis.xyz is not None
-            and j.axis.xyz.xyz is not None
-            else None,
+            axis=(
+                np.array(j.axis.xyz.xyz)
+                if j.axis is not None
+                and j.axis.xyz is not None
+                and j.axis.xyz.xyz is not None
+                else None
+            ),
             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.upper)
-                if j.axis is not None and j.axis.limit is not None
-                else np.finfo(float).max,
+                (
+                    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.upper)
+                    if j.axis is not None and j.axis.limit is not None
+                    else np.finfo(float).max
+                ),
+            ),
+            friction_static=(
+                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=(
+                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=(
+                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=(
+                j.axis.limit.stiffness
+                if j.axis is not None
+                and j.axis.limit is not None
+                and j.axis.limit.stiffness is not None
+                else 0.0
             ),
-            friction_static=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=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=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=j.axis.limit.stiffness
-            if j.axis is not None
-            and j.axis.limit is not None
-            and j.axis.limit.stiffness is not None
-            else 0.0,
         )
         for j in sdf_model.joints()
         if j.type in {"revolute", "prismatic", "fixed"}
@@ -341,6 +357,6 @@ def build_model_description(
     )
 
     # Store the parsed SDF tree as extra info
-    model = dataclasses.replace(model, extra_info=dict(sdf_model=sdf_data.sdf_model))
+    model = dataclasses.replace(model, extra_info={"sdf_model": sdf_data.sdf_model})
 
     return model

jaxsim/parsers/rod/utils.py

@@ -1,15 +1,17 @@
 import os
 from typing import Union
 
-import jax.numpy as jnp
+import jaxlie
 import numpy as np
 import numpy.typing as npt
 import rod
 
+import jaxsim.typing as jtp
+from jaxsim.math.inertia import Inertia
 from jaxsim.parsers import descriptions
 
 
-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,9 +22,6 @@ 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
     m = inertial.mass
 
@@ -52,13 +51,13 @@ def from_sdf_inertial(inertial: rod.Inertial) -> npt.NDArray:
     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()
+    CoM_H_L = jaxlie.SE3.from_matrix(matrix=L_H_CoM).inverse()
+    CoM_X_L = CoM_H_L.adjoint()
 
     # 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(

jaxsim/rbda/__init__.py

@@ -0,0 +1,7 @@
+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_full_doubly_left
+from .rnea import rnea
+from .soft_contacts import SoftContacts, SoftContactsParams

jaxsim/rbda/aba.py

@@ -0,0 +1,295 @@
+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 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.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))
+    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.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
+    ]
+
+    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.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
+    ]
+
+    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.MatrixJax, jtp.MatrixJax]
+        ) -> tuple[jtp.MatrixJax, jtp.MatrixJax]:
+
+            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.MatrixJax, jtp.VectorJax]
+    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())

jaxsim/rbda/collidable_points.py

@@ -0,0 +1,142 @@
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Quaternion, Skew
+
+from . import utils
+
+
+def collidable_points_pos_vel(
+    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,
+) -> tuple[jtp.Matrix, jtp.Matrix]:
+    """
+
+    Compute the position and linear velocity of collidable points in the world frame.
+
+    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.
+
+    Returns:
+        A tuple containing the position and linear velocity of collidable points.
+    """
+
+    if len(model.kin_dyn_parameters.contact_parameters.body) == 0:
+        return jnp.array(0).astype(float), jnp.empty(0).astype(float)
+
+    W_p_B, W_Q_B, s, W_v_WB, ṡ, _, _, _, _, _ = 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,
+    )
+
+    # 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 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 buffer of transforms world -> link and initialize the base pose.
+    W_X_i = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    W_X_i = W_X_i.at[0].set(Adjoint.inverse(i_X_λi[0]))
+
+    # Allocate buffer of 6D inertial-fixed velocities and initialize the base velocity.
+    W_v_Wi = jnp.zeros(shape=(model.number_of_links(), 6))
+    W_v_Wi = W_v_Wi.at[0].set(W_v_WB)
+
+    # ====================
+    # Propagate kinematics
+    # ====================
+
+    PropagateTransformsCarry = tuple[jtp.MatrixJax, jtp.Matrix]
+    propagate_transforms_carry: PropagateTransformsCarry = (W_X_i, W_v_Wi)
+
+    def propagate_kinematics(
+        carry: PropagateTransformsCarry, i: jtp.Int
+    ) -> tuple[PropagateTransformsCarry, None]:
+
+        ii = i - 1
+        W_X_i, W_v_Wi = carry
+
+        # Compute the parent to child 6D transform.
+        λi_X_i = Adjoint.inverse(adjoint=i_X_λi[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)
+
+        # Propagate the 6D velocity
+        W_vi_Wi = W_v_Wi[λ[i]] + W_X_i[i] @ (S[i] * ṡ[ii]).squeeze()
+        W_v_Wi = W_v_Wi.at[i].set(W_vi_Wi)
+
+        return (W_X_i, W_v_Wi), None
+
+    (W_X_i, W_v_Wi), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_transforms_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(W_X_i, W_v_Wi), None]
+    )
+
+    # ==================================================
+    # 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_vl_WC = jax.vmap(process_point_kinematics)(
+        model.kin_dyn_parameters.contact_parameters.point,
+        jnp.array(model.kin_dyn_parameters.contact_parameters.body),
+    )
+
+    return W_p_Ci, CW_vl_WC