jaxsim 0.2.dev388__py3-none-any.whl → 0.2.dev394__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.dev388'
-__version_tuple__ = version_tuple = (0, 2, 'dev388')
+__version__ = version = '0.2.dev394'
+__version_tuple__ = version_tuple = (0, 2, 'dev394')

jaxsim/api/link.py

@@ -335,3 +335,27 @@ def velocity(
 
     # Compute the link velocity in the output velocity representation.
     return O_J_WL_I @ I_ν
+
+
+@jax.jit
+def bias_acceleration(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_index: jtp.IntLike,
+) -> jtp.Vector:
+    """
+    Compute the bias acceleration of the link.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_index: The index of the link.
+
+    Returns:
+        The 6D bias acceleration of the link.
+    """
+
+    # Compute the bias acceleration of all links in the active representation.
+    O_v̇_WL = js.model.link_bias_accelerations(model=model, data=data)[link_index]
+    return O_v̇_WL

jaxsim/api/model.py

@@ -1243,6 +1243,211 @@ def average_velocity_jacobian(
 # ========================
 
 
+@jax.jit
+def link_bias_accelerations(
+    model: JaxSimModel,
+    data: js.data.JaxSimModelData,
+) -> jtp.Vector:
+    r"""
+    Compute the bias accelerations of the links of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The bias accelerations of the links of the model.
+
+    Note:
+        This function computes the component of the total 6D acceleration not due to
+        the joint or base acceleration.
+        It is often called :math:`\dot{J} \boldsymbol{\nu}`.
+    """
+
+    # ================================================
+    # Compute the body-fixed zero base 6D acceleration
+    # ================================================
+
+    # Compute the base transform.
+    W_H_B = jaxlie.SE3.from_rotation_and_translation(
+        rotation=jaxlie.SO3.from_quaternion_xyzw(
+            xyzw=jaxsim.math.Quaternion.to_xyzw(wxyz=data.base_orientation())
+        ),
+        translation=data.base_position(),
+    ).as_matrix()
+
+    def other_representation_to_inertial(
+        C_v̇_WB: jtp.Vector, C_v_WB: jtp.Vector, W_H_C: jtp.Matrix, W_v_WC: jtp.Vector
+    ) -> jtp.Vector:
+        """
+        Helper to convert the active representation of the base acceleration C_v̇_WB
+        expressed in a generic frame C to the inertial-fixed representation W_v̇_WB.
+        """
+
+        W_X_C = jaxlie.SE3.from_matrix(W_H_C).adjoint()
+        C_X_W = jaxlie.SE3.from_matrix(W_H_C).inverse().adjoint()
+
+        # In Mixed representation, we need to include a cross product in ℝ⁶.
+        # In Inertial and Body representations, the cross product is always zero.
+        return W_X_C @ (C_v̇_WB + jaxsim.math.Cross.vx(C_X_W @ W_v_WC) @ C_v_WB)
+
+    # Here we initialize a zero 6D acceleration in the active representation, and
+    # convert it to inertial-fixed. This is a useful intermediate representation
+    # because the apparent acceleration W_v̇_WB is equal to the intrinsic acceleration
+    # W_a_WB, and intrinsic accelerations can be expressed in different frames through
+    # a simple C_X_W 6D transform.
+    match data.velocity_representation:
+        case VelRepr.Inertial:
+            W_H_C = W_H_W = jnp.eye(4)
+            W_v_WC = W_v_WW = jnp.zeros(6)
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                C_v_WB = W_v_WB = data.base_velocity()
+
+        case VelRepr.Body:
+            W_H_C = W_H_B
+            with data.switch_velocity_representation(VelRepr.Inertial):
+                W_v_WC = W_v_WB = data.base_velocity()
+            with data.switch_velocity_representation(VelRepr.Body):
+                C_v_WB = B_v_WB = data.base_velocity()
+
+        case VelRepr.Mixed:
+            W_H_BW = W_H_B.at[0:3, 0:3].set(jnp.eye(3))
+            W_H_C = W_H_BW
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                W_ṗ_B = data.base_velocity()[0:3]
+                W_v_WC = W_v_W_BW = jnp.zeros(6).at[0:3].set(W_ṗ_B)
+            with data.switch_velocity_representation(VelRepr.Mixed):
+                C_v_WB = BW_v_WB = data.base_velocity()
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Convert a zero 6D acceleration from the active representation to inertial-fixed.
+    W_v̇_WB = other_representation_to_inertial(
+        C_v̇_WB=jnp.zeros(6), C_v_WB=C_v_WB, W_H_C=W_H_C, W_v_WC=W_v_WC
+    )
+
+    # ===================================
+    # Initialize buffers and prepare data
+    # ===================================
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute 6D transforms of the base velocity.
+    B_X_W = jaxsim.math.Adjoint.from_transform(transform=W_H_B, inverse=True)
+
+    # 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=data.joint_positions(), base_transform=W_H_B
+    )
+
+    # Allocate the buffer to store the body-fixed link velocities.
+    L_v_WL = jnp.zeros(shape=(model.number_of_links(), 6))
+
+    # Store the base velocity.
+    with data.switch_velocity_representation(VelRepr.Body):
+        B_v_WB = data.base_velocity()
+        L_v_WL = L_v_WL.at[0].set(B_v_WB)
+
+    # Get the joint velocities.
+    ṡ = data.joint_velocities(model=model, joint_names=model.joint_names())
+
+    # Allocate the buffer to store the body-fixed link accelerations,
+    # and initialize the base acceleration.
+    L_v̇_WL = jnp.zeros(shape=(model.number_of_links(), 6))
+    L_v̇_WL = L_v̇_WL.at[0].set(B_X_W @ W_v̇_WB)
+
+    # ======================================
+    # Propagate accelerations and velocities
+    # ======================================
+
+    # The computation of the bias forces is similar to the forward pass of RNEA,
+    # this time with zero base and joint accelerations. Furthermore, here we do
+    # not remove gravity during the propagation.
+
+    # Initialize the loop.
+    Carry = tuple[jtp.MatrixJax, jtp.MatrixJax]
+    carry0: Carry = (L_v_WL, L_v̇_WL)
+
+    def propagate_accelerations(carry: Carry, i: jtp.Int) -> tuple[Carry, None]:
+        # Initialize index and unpack the carry.
+        ii = i - 1
+        v, a = carry
+
+        # Get the motion subspace of the joint.
+        Si = S[i].squeeze()
+
+        # Project the joint velocity into its motion subspace.
+        vJ = Si * ṡ[ii]
+
+        # Propagate the link body-fixed velocity.
+        v_i = i_X_λi[i] @ v[λ[i]] + vJ
+        v = v.at[i].set(v_i)
+
+        # Propagate the link body-fixed acceleration considering zero joint acceleration.
+        s̈ = 0.0
+        a_i = i_X_λi[i] @ a[λ[i]] + Si * s̈ + jaxsim.math.Cross.vx(v[i]) @ vJ
+        a = a.at[i].set(a_i)
+
+        return (v, a), None
+
+    # Compute the body-fixed velocity and body-fixed apparent acceleration of the links.
+    (L_v_WL, L_v̇_WL), _ = (
+        jax.lax.scan(
+            f=propagate_accelerations,
+            init=carry0,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(L_v_WL, L_v̇_WL), None]
+    )
+
+    # ===================================================================
+    # Convert the body-fixed 6D acceleration to the active representation
+    # ===================================================================
+
+    def body_to_other_representation(
+        L_v̇_WL: jtp.Vector, L_v_WL: jtp.Vector, C_H_L: jtp.Matrix, L_v_CL: jtp.Vector
+    ) -> jtp.Vector:
+        """
+        Helper to convert the body-fixed apparent acceleration L_v̇_WL to
+        another representation C_v̇_WL expressed in a generic frame C.
+        """
+
+        # In Mixed representation, we need to include a cross product in ℝ⁶.
+        # In Inertial and Body representations, the cross product is always zero.
+        C_X_L = jaxsim.math.Adjoint.from_transform(transform=C_H_L)
+        return C_X_L @ (L_v̇_WL + jaxsim.math.Cross.vx(L_v_CL) @ L_v_WL)
+
+    match data.velocity_representation:
+        case VelRepr.Body:
+            C_H_L = L_H_L = jnp.stack([jnp.eye(4)] * model.number_of_links())
+            L_v_CL = L_v_LL = jnp.zeros(shape=(model.number_of_links(), 6))
+
+        case VelRepr.Inertial:
+            C_H_L = W_H_L = js.model.forward_kinematics(model=model, data=data)
+            L_v_CL = L_v_WL
+
+        case VelRepr.Mixed:
+            W_H_L = js.model.forward_kinematics(model=model, data=data)
+            LW_H_L = jax.vmap(lambda W_H_L: W_H_L.at[0:3, 3].set(jnp.zeros(3)))(W_H_L)
+            C_H_L = LW_H_L
+            L_v_CL = L_v_LW_L = jax.vmap(lambda v: v.at[0:3].set(jnp.zeros(3)))(L_v_WL)
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Convert from body-fixed to the active representation.
+    O_v̇_WL = jax.vmap(body_to_other_representation)(
+        L_v̇_WL=L_v̇_WL, L_v_WL=L_v_WL, C_H_L=C_H_L, L_v_CL=L_v_CL
+    )
+
+    return O_v̇_WL
+
+
 @jax.jit
 def link_contact_forces(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData

{jaxsim-0.2.dev388.dist-info → jaxsim-0.2.dev394.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.dev388
+Version: 0.2.dev394
 Summary: A physics engine in reduced coordinates implemented with JAX.
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo

{jaxsim-0.2.dev388.dist-info → jaxsim-0.2.dev394.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=OcrfoYS1DGcmAGqu2AqlCTiUVxcpi-IsVwcr_16x74Q,1789
-jaxsim/_version.py,sha256=ajN4RtI-FZYlNpUqVy327D57EXvmjZkyMCV_lzLqeUo,423
+jaxsim/_version.py,sha256=so5jZFKl8aDmZRylEyYqJKNNSSEg2aGlO-xfDv0En8w,423
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=MeuOCQtLAr-sPkvB_sU8FtwGNRirz1auCwIgRC-QZl8,646
 jaxsim/api/__init__.py,sha256=fNTCPUeDfOAizRd4RsW3Epv0sLTu0KJGoFRSEsi75VM,162
@@ -9,8 +9,8 @@ jaxsim/api/contact.py,sha256=Ve4ZOWkLEBRgK3KhtICxKY7YzsxYvc3lO-pPRBjqSnY,8659
 jaxsim/api/data.py,sha256=cPFy6TTm0rNpDk1O0ZxOpvBwWVPu6V14XLRMeC7x81Q,26786
 jaxsim/api/joint.py,sha256=q31Kp3Cqv-yTcxijjzbj_QADFnGQyjb2al9fYZtzedo,4763
 jaxsim/api/kin_dyn_parameters.py,sha256=G4mtSi8fElYe0yttLgsxSOPf7vcK-yqTu06Aa5SSrYg,26012
-jaxsim/api/link.py,sha256=LZVcQhQsTKsfR13KewFtEMYu4siVJl7mqoDwYsoFFes,9240
-jaxsim/api/model.py,sha256=Kpe5PQHP1VCS8Xmu4M_aa6EAJFPkOCO_Y11-TGhEtq8,44091
+jaxsim/api/link.py,sha256=rypTwkMf9HJ5UuAtHRJh0LqqdJWcLKTtTjWcjduEsF0,9842
+jaxsim/api/model.py,sha256=u7CZTNc0UY0ZVpUEW-RsZi2UysyzbK06zvSPsdOlPYQ,52269
 jaxsim/api/ode.py,sha256=6l-6i2YHagsQvR8Ac-_fmO6P0hBVT6NkHhwXnrdITEg,9785
 jaxsim/api/ode_data.py,sha256=dwRFVPJ30XMmdUbPXEu7YxsQ97jZP4L4fd5ZzhrO5Ys,22184
 jaxsim/api/references.py,sha256=Lvskf17r619KKxwCJP7hAAty2kaXgDXJX1uKqoDIDgo,15483
@@ -57,8 +57,8 @@ jaxsim/utils/__init__.py,sha256=tnQq1_CavdfeKaLYt3pmO7Jk4MU2RwwQU_qICkjyoTY,197
 jaxsim/utils/hashless.py,sha256=bFIwKeo9KiWwsY8QM55duEGGQOyyJ4jQyPcuqTLEp5k,297
 jaxsim/utils/jaxsim_dataclass.py,sha256=T8nZolb563mYJD7YvTmYW-z-hpkEiaK3cDbYHqwtqRc,11347
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim-0.2.dev388.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
-jaxsim-0.2.dev388.dist-info/METADATA,sha256=T1FlCTdSjlgSIOFqv9Fik7fMwpwIcX6Ps5N2fCElQ7I,7630
-jaxsim-0.2.dev388.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.2.dev388.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.2.dev388.dist-info/RECORD,,
+jaxsim-0.2.dev394.dist-info/LICENSE,sha256=EsU2z6_sWW4Zduzq3goVWjZoCZVKQsM4H_y0o7oRA7Q,1547
+jaxsim-0.2.dev394.dist-info/METADATA,sha256=MAzjWBb9UwRC3HkwkBP6ivGKFF7qKmbUT0oGDyB0VJE,7630
+jaxsim-0.2.dev394.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.2.dev394.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.2.dev394.dist-info/RECORD,,