jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev5__py3-none-any.whl

jaxsim/math/quaternion.py

@@ -1,21 +1,27 @@
 import jax.lax
 import jax.numpy as jnp
+import jaxlie
 
 import jaxsim.typing as jtp
-from jaxsim.sixd import so3
+
+from .utils import safe_norm
 
 
 class Quaternion:
+    """
+    A utility class for quaternion operations.
+    """
+
     @staticmethod
     def to_xyzw(wxyz: jtp.Vector) -> jtp.Vector:
         """
         Convert a quaternion from WXYZ to XYZW representation.
 
         Args:
-            wxyz (jtp.Vector): Quaternion in WXYZ representation.
+            wxyz: Quaternion in WXYZ representation.
 
         Returns:
-            jtp.Vector: Quaternion in XYZW representation.
+            Quaternion in XYZW representation.
         """
         return wxyz.squeeze()[jnp.array([1, 2, 3, 0])]
 
@@ -25,10 +31,10 @@ class Quaternion:
         Convert a quaternion from XYZW to WXYZ representation.
 
         Args:
-            xyzw (jtp.Vector): Quaternion in XYZW representation.
+            xyzw: Quaternion in XYZW representation.
 
         Returns:
-            jtp.Vector: Quaternion in WXYZ representation.
+            Quaternion in WXYZ representation.
         """
         return xyzw.squeeze()[jnp.array([3, 0, 1, 2])]
 
@@ -38,14 +44,12 @@ class Quaternion:
         Convert a quaternion to a direction cosine matrix (DCM).
 
         Args:
-            quaternion (jtp.Vector): Quaternion in XYZW representation.
+            quaternion: Quaternion in XYZW representation.
 
         Returns:
-            jtp.Matrix: Direction cosine matrix (DCM).
+            The Direction cosine matrix (DCM).
         """
-        return so3.SO3.from_quaternion_xyzw(
-            xyzw=Quaternion.to_xyzw(quaternion)
-        ).as_matrix()
+        return jaxlie.SO3(wxyz=quaternion).as_matrix()
 
     @staticmethod
     def from_dcm(dcm: jtp.Matrix) -> jtp.Vector:
@@ -53,14 +57,12 @@ class Quaternion:
         Convert a direction cosine matrix (DCM) to a quaternion.
 
         Args:
-            dcm (jtp.Matrix): Direction cosine matrix (DCM).
+            dcm: Direction cosine matrix (DCM).
 
         Returns:
-            jtp.Vector: Quaternion in XYZW representation.
+            Quaternion in WXYZ representation.
         """
-        return Quaternion.to_wxyz(
-            xyzw=so3.SO3.from_matrix(matrix=dcm).as_quaternion_xyzw()
-        )
+        return jaxlie.SO3.from_matrix(matrix=dcm).wxyz
 
     @staticmethod
     def derivative(
@@ -73,13 +75,13 @@ class Quaternion:
         Compute the derivative of a quaternion given angular velocity.
 
         Args:
-            quaternion (jtp.Vector): Quaternion in XYZW representation.
-            omega (jtp.Vector): Angular velocity vector.
+            quaternion: Quaternion in XYZW representation.
+            omega: Angular velocity vector.
             omega_in_body_fixed (bool): Whether the angular velocity is in the body-fixed frame.
             K (float): A scaling factor.
 
         Returns:
-            jtp.Vector: The derivative of the quaternion.
+            The derivative of the quaternion.
         """
         ω = omega.squeeze()
         quaternion = quaternion.squeeze()
@@ -115,21 +117,53 @@ class Quaternion:
             operand=quaternion,
         )
 
-        norm_ω = jax.lax.cond(
-            pred=ω.dot(ω) < (1e-6) ** 2,
-            true_fun=lambda _: 1e-6,
-            false_fun=lambda _: jnp.linalg.norm(ω),
-            operand=None,
-        )
+        norm_ω = safe_norm(ω)
 
         qd = 0.5 * (
             Q
             @ jnp.hstack(
                 [
-                    K * norm_ω * (1 - jnp.linalg.norm(quaternion)),
+                    K * norm_ω * (1 - safe_norm(quaternion)),
                     ω,
                 ]
             )
         )
 
         return jnp.vstack(qd)
+
+    @staticmethod
+    def integration(
+        quaternion: jtp.VectorLike,
+        dt: jtp.FloatLike,
+        omega: jtp.VectorLike,
+        omega_in_body_fixed: jtp.BoolLike = False,
+    ) -> jtp.Vector:
+        """
+        Integrate a quaternion in SO(3) given an angular velocity.
+
+        Args:
+            quaternion: The quaternion to integrate.
+            dt: The time step.
+            omega: The angular velocity vector.
+            omega_in_body_fixed:
+                Whether the angular velocity is in body-fixed representation
+                as opposed to the default inertial-fixed representation.
+
+        Returns:
+            The integrated quaternion.
+        """
+
+        ω_AB = jnp.array(omega).squeeze().astype(float)
+        A_Q_B = jnp.array(quaternion).squeeze().astype(float)
+
+        # Build the initial SO(3) quaternion.
+        W_Q_B_t0 = jaxlie.SO3(wxyz=A_Q_B)
+
+        # Integrate the quaternion on the manifold.
+        W_Q_B_tf = jax.lax.select(
+            pred=omega_in_body_fixed,
+            on_true=(W_Q_B_t0 @ jaxlie.SO3.exp(tangent=dt * ω_AB)).wxyz,
+            on_false=(jaxlie.SO3.exp(tangent=dt * ω_AB) @ W_Q_B_t0).wxyz,
+        )
+
+        return W_Q_B_tf

jaxsim/math/rotation.py

@@ -1,27 +1,30 @@
-from typing import Tuple
-
-import jax
 import jax.numpy as jnp
+import jaxlie
 
 import jaxsim.typing as jtp
-from jaxsim.sixd import so3
 
 from .skew import Skew
+from .utils import safe_norm
 
 
 class Rotation:
+    """
+    A utility class for rotation matrix operations.
+    """
+
     @staticmethod
     def x(theta: jtp.Float) -> jtp.Matrix:
         """
         Generate a 3D rotation matrix around the X-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
-            jtp.Matrix: 3D rotation matrix.
+            The 3D rotation matrix.
         """
-        return so3.SO3.from_x_radians(theta=theta).as_matrix()
+
+        return jaxlie.SO3.from_x_radians(theta=theta).as_matrix()
 
     @staticmethod
     def y(theta: jtp.Float) -> jtp.Matrix:
@@ -29,12 +32,13 @@ class Rotation:
         Generate a 3D rotation matrix around the Y-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
-            jtp.Matrix: 3D rotation matrix.
+            The 3D rotation matrix.
         """
-        return so3.SO3.from_y_radians(theta=theta).as_matrix()
+
+        return jaxlie.SO3.from_y_radians(theta=theta).as_matrix()
 
     @staticmethod
     def z(theta: jtp.Float) -> jtp.Matrix:
@@ -42,12 +46,13 @@ class Rotation:
         Generate a 3D rotation matrix around the Z-axis.
 
         Args:
-            theta (jtp.Float): Rotation angle in radians.
+            theta: Rotation angle in radians.
 
         Returns:
-            jtp.Matrix: 3D rotation matrix.
+            The 3D rotation matrix.
         """
-        return so3.SO3.from_z_radians(theta=theta).as_matrix()
+
+        return jaxlie.SO3.from_z_radians(theta=theta).as_matrix()
 
     @staticmethod
     def from_axis_angle(vector: jtp.Vector) -> jtp.Matrix:
@@ -55,17 +60,18 @@ class Rotation:
         Generate a 3D rotation matrix from an axis-angle representation.
 
         Args:
-            vector (jtp.Vector): Axis-angle representation as a 3D vector.
+            vector: Axis-angle representation or the rotation as a 3D vector.
 
         Returns:
-            jtp.Matrix: 3D rotation matrix.
-
+            The SO(3) rotation matrix.
         """
+
         vector = vector.squeeze()
-        theta = jnp.linalg.norm(vector)
 
-        def theta_is_not_zero(theta_and_v: Tuple[jtp.Float, jtp.Vector]) -> jtp.Matrix:
-            theta, v = theta_and_v
+        def theta_is_not_zero(axis: jtp.Vector) -> jtp.Matrix:
+
+            v = axis
+            theta = safe_norm(v)
 
             s = jnp.sin(theta)
             c = jnp.cos(theta)
@@ -79,9 +85,9 @@ class Rotation:
 
             return R.transpose()
 
-        return jax.lax.cond(
-            pred=(theta == 0.0),
-            true_fun=lambda operand: jnp.eye(3),
-            false_fun=theta_is_not_zero,
-            operand=(theta, vector),
+        return jnp.where(
+            jnp.allclose(vector, 0.0),
+            # Return an identity rotation matrix when the input vector is zero.
+            jnp.eye(3),
+            theta_is_not_zero(axis=vector),
         )

jaxsim/math/skew.py

@@ -14,15 +14,26 @@ class Skew:
         Compute the skew-symmetric matrix (wedge operator) of a 3D vector.
 
         Args:
-            vector (jtp.Vector): A 3D vector.
+            vector: A 3D vector.
 
         Returns:
-            jtp.Matrix: The skew-symmetric matrix corresponding to the input vector.
+            The skew-symmetric matrix corresponding to the input vector.
 
         """
-        vector = vector.squeeze()
-        x, y, z = vector
-        skew = jnp.array([[0, -z, y], [z, 0, -x], [-y, x, 0]])
+
+        vector = vector.reshape(-1, 3)
+
+        x, y, z = jnp.split(vector, 3, axis=-1)
+
+        skew = jnp.stack(
+            [
+                jnp.concatenate([jnp.zeros_like(x), -z, y], axis=-1),
+                jnp.concatenate([z, jnp.zeros_like(x), -x], axis=-1),
+                jnp.concatenate([-y, x, jnp.zeros_like(x)], axis=-1),
+            ],
+            axis=-2,
+        ).squeeze()
+
         return skew
 
     @staticmethod
@@ -31,10 +42,10 @@ class Skew:
         Extract the 3D vector from a skew-symmetric matrix (vee operator).
 
         Args:
-            matrix (jtp.Matrix): A 3x3 skew-symmetric matrix.
+            matrix: A 3x3 skew-symmetric matrix.
 
         Returns:
-            jtp.Vector: The 3D vector extracted from the input matrix.
+            The 3D vector extracted from the input matrix.
 
         """
         vector = 0.5 * jnp.vstack(

jaxsim/math/transform.py

@@ -0,0 +1,102 @@
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.typing as jtp
+
+
+class Transform:
+    """
+    A utility class for transformation matrix operations.
+    """
+
+    @staticmethod
+    def from_quaternion_and_translation(
+        quaternion: jtp.VectorLike | None = None,
+        translation: jtp.VectorLike | None = None,
+        inverse: jtp.BoolLike = False,
+        normalize_quaternion: jtp.BoolLike = False,
+    ) -> jtp.Matrix:
+        """
+        Create a transformation matrix from a quaternion and a translation.
+
+        Args:
+            quaternion: A 4D vector representing a SO(3) orientation.
+            translation: A 3D vector representing a translation.
+            inverse: Whether to compute the inverse transformation.
+            normalize_quaternion:
+                Whether to normalize the quaternion before creating the transformation.
+
+        Returns:
+            The 4x4 transformation matrix representing the SE(3) transformation.
+        """
+
+        quaternion = quaternion if quaternion is not None else jnp.array([1.0, 0, 0, 0])
+        translation = translation if translation is not None else jnp.zeros(3)
+
+        W_Q_B = jnp.array(quaternion).astype(float)
+        W_p_B = jnp.array(translation).astype(float)
+
+        assert W_p_B.size == 3
+        assert W_Q_B.size == 4
+
+        A_R_B = jaxlie.SO3(wxyz=W_Q_B)
+        A_R_B = A_R_B if not normalize_quaternion else A_R_B.normalize()
+
+        A_H_B = jaxlie.SE3.from_rotation_and_translation(
+            rotation=A_R_B, translation=W_p_B
+        )
+
+        return A_H_B.as_matrix() if not inverse else A_H_B.inverse().as_matrix()
+
+    @staticmethod
+    def from_rotation_and_translation(
+        rotation: jtp.MatrixLike | None = None,
+        translation: jtp.VectorLike | None = None,
+        inverse: jtp.BoolLike = False,
+    ) -> jtp.Matrix:
+        """
+        Create a transformation matrix from a rotation matrix and a translation vector.
+
+        Args:
+            rotation: A 3x3 rotation matrix representing a SO(3) orientation.
+            translation: A 3D vector representing a translation.
+            inverse: Whether to compute the inverse transformation.
+
+        Returns:
+            The 4x4 transformation matrix representing the SE(3) transformation.
+        """
+        rotation = rotation if rotation is not None else jnp.eye(3)
+        translation = translation if translation is not None else jnp.zeros(3)
+
+        A_R_B = jnp.array(rotation).astype(float)
+        W_p_B = jnp.array(translation).astype(float)
+
+        assert W_p_B.size == 3
+        assert A_R_B.shape == (3, 3)
+
+        A_H_B = jaxlie.SE3.from_rotation_and_translation(
+            rotation=jaxlie.SO3.from_matrix(A_R_B), translation=W_p_B
+        )
+
+        return A_H_B.as_matrix() if not inverse else A_H_B.inverse().as_matrix()
+
+    @staticmethod
+    def inverse(transform: jtp.MatrixLike) -> jtp.Matrix:
+        """
+        Compute the inverse transformation matrix.
+
+        Args:
+            transform: A 4x4 transformation matrix.
+
+        Returns:
+            The 4x4 inverse transformation matrix.
+        """
+
+        A_H_B = jnp.reshape(transform, (-1, 4, 4))
+
+        return (
+            jaxlie.SE3.from_matrix(matrix=A_H_B)
+            .inverse()
+            .as_matrix()
+            .reshape(transform.shape[:-2] + (4, 4))
+        )

jaxsim/math/utils.py

@@ -0,0 +1,31 @@
+import jax.numpy as jnp
+
+import jaxsim.typing as jtp
+
+
+def safe_norm(array: jtp.ArrayLike, axis=None) -> jtp.Array:
+    """
+    Compute an array norm handling NaNs and making sure that
+    it is safe to get the gradient.
+
+    Args:
+        array: The array for which to compute the norm.
+        axis: The axis for which to compute the norm.
+
+    Returns:
+        The norm of the array with handling for zero arrays to avoid NaNs.
+    """
+
+    # Check if the entire array is composed of zeros.
+    is_zero = jnp.allclose(array, 0.0)
+
+    # Replace zeros with an array of ones temporarily to avoid division by zero.
+    # This ensures the computation of norm does not produce NaNs or Infs.
+    array = jnp.where(is_zero, jnp.ones_like(array), array)
+
+    # Compute the norm of the array along the specified axis.
+    norm = jnp.linalg.norm(array, axis=axis)
+
+    # Use `jnp.where` to set the norm to 0.0 where the input array was all zeros.
+    # This usage supports potential batch processing for future scalability.
+    return jnp.where(is_zero, 0.0, norm)

jaxsim/mujoco/__init__.py

@@ -1,3 +1,4 @@
-from .loaders import RodModelToMjcf, SdfToMjcf, UrdfToMjcf
+from .loaders import ModelToMjcf, RodModelToMjcf, SdfToMjcf, UrdfToMjcf
 from .model import MujocoModelHelper
+from .utils import mujoco_data_from_jaxsim
 from .visualizer import MujocoVideoRecorder, MujocoVisualizer