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

jaxsim/__init__.py

@@ -6,12 +6,12 @@ from ._version import __version__
 def _jnp_options() -> None:
     import os
 
-    from jax.config import config
+    import jax
 
     # Enable by default
     if not ("JAX_ENABLE_X64" in os.environ and os.environ["JAX_ENABLE_X64"] == "0"):
         logging.info("Enabling JAX to use 64bit precision")
-        config.update("jax_enable_x64", True)
+        jax.config.update("jax_enable_x64", True)
 
         import jax.numpy as jnp
         import numpy as np
@@ -61,7 +61,6 @@ del _jnp_options
 del _np_options
 del _is_editable
 
-from . import high_level, logging, math, simulation, sixd
-from .high_level.common import VelRepr
-from .simulation.ode_integration import IntegratorType
-from .simulation.simulator import JaxSim
+from . import terrain  # isort:skip
+from . import api, integrators, logging, math, rbda
+from .api.common import VelRepr

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1rc0'
-__version_tuple__ = version_tuple = (0, 1)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

jaxsim/api/__init__.py

@@ -0,0 +1,3 @@
+from . import common  # isort:skip
+from . import model, data  # isort:skip
+from . import com, contact, joint, kin_dyn_parameters, link, ode, ode_data, references

jaxsim/api/com.py

@@ -0,0 +1,240 @@
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.math
+import jaxsim.typing as jtp
+
+from .common import VelRepr
+
+
+@jax.jit
+def com_position(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    """
+    Compute the position of the center of mass of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The position of the center of mass of the model w.r.t. the world frame.
+    """
+
+    m = js.model.total_mass(model=model)
+
+    W_H_L = js.model.forward_kinematics(model=model, data=data)
+    W_H_B = data.base_transform()
+    B_H_W = jaxlie.SE3.from_matrix(W_H_B).inverse().as_matrix()
+
+    def B_p̃_LCoM(i) -> jtp.Vector:
+        m = js.link.mass(model=model, link_index=i)
+        L_p_LCoM = js.link.com_position(
+            model=model, data=data, link_index=i, in_link_frame=True
+        )
+        return m * B_H_W @ W_H_L[i] @ jnp.hstack([L_p_LCoM, 1])
+
+    com_links = jax.vmap(B_p̃_LCoM)(jnp.arange(model.number_of_links()))
+
+    B_p̃_CoM = (1 / m) * com_links.sum(axis=0)
+    B_p̃_CoM = B_p̃_CoM.at[3].set(1)
+
+    return (W_H_B @ B_p̃_CoM)[0:3].astype(float)
+
+
+@jax.jit
+def com_linear_velocity(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the linear velocity of the center of mass of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The linear velocity of the center of mass of the model in the
+        active representation.
+
+    Note:
+        The linear velocity of the center of mass  is expressed in the mixed frame
+        :math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
+        active velocity representation is either inertial-fixed or mixed,
+        and :math:`[C] = [B]` if the active velocity representation is body-fixed.
+    """
+
+    # Extract the linear component of the 6D average centroidal velocity.
+    # This is expressed in G[B] in body-fixed representation, and in G[W] in
+    # inertial-fixed or mixed representation.
+    G_vl_WG = average_centroidal_velocity(model=model, data=data)[0:3]
+
+    return G_vl_WG
+
+
+@jax.jit
+def centroidal_momentum(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the centroidal momentum of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The centroidal momentum of the model.
+
+    Note:
+        The centroidal momentum is expressed in the mixed frame
+        :math:`({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`C = W` if the
+        active velocity representation is either inertial-fixed or mixed,
+        and :math:`C = B` if the active velocity representation is body-fixed.
+    """
+
+    ν = data.generalized_velocity()
+    G_J = centroidal_momentum_jacobian(model=model, data=data)
+
+    return G_J @ ν
+
+
+@jax.jit
+def centroidal_momentum_jacobian(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    r"""
+    Compute the Jacobian of the centroidal momentum of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The Jacobian of the centroidal momentum of the model.
+
+    Note:
+        The frame corresponding to the output representation of this Jacobian is either
+        :math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
+        or :math:`G[B]`, if the active velocity representation is body-fixed.
+
+    Note:
+        This Jacobian is also known in the literature as Centroidal Momentum Matrix.
+    """
+
+    # Compute the Jacobian of the total momentum with body-fixed output representation.
+    # We convert the output representation either to G[W] or G[B] below.
+    B_Jh = js.model.total_momentum_jacobian(
+        model=model, data=data, output_vel_repr=VelRepr.Body
+    )
+
+    W_H_B = data.base_transform()
+    B_H_W = jaxsim.math.Transform.inverse(W_H_B)
+
+    W_p_CoM = com_position(model=model, data=data)
+
+    match data.velocity_representation:
+        case VelRepr.Inertial | VelRepr.Mixed:
+            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
+        case VelRepr.Body:
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Compute the transform for 6D forces.
+    G_Xf_B = jaxsim.math.Adjoint.from_transform(transform=B_H_W @ W_H_G).T
+
+    return G_Xf_B @ B_Jh
+
+
+@jax.jit
+def locked_centroidal_spatial_inertia(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+):
+    """
+    Compute the locked centroidal spatial inertia of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The locked centroidal spatial inertia of the model.
+    """
+
+    with data.switch_velocity_representation(VelRepr.Body):
+        B_Mbb_B = js.model.locked_spatial_inertia(model=model, data=data)
+
+    W_H_B = data.base_transform()
+    W_p_CoM = com_position(model=model, data=data)
+
+    match data.velocity_representation:
+        case VelRepr.Inertial | VelRepr.Mixed:
+            W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
+        case VelRepr.Body:
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    B_H_G = jaxlie.SE3.from_matrix(jaxsim.math.Transform.inverse(W_H_B) @ W_H_G)
+
+    B_Xv_G = B_H_G.adjoint()
+    G_Xf_B = B_Xv_G.transpose()
+
+    return G_Xf_B @ B_Mbb_B @ B_Xv_G
+
+
+@jax.jit
+def average_centroidal_velocity(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the average centroidal velocity of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The average centroidal velocity of the model.
+
+    Note:
+        The average velocity is expressed in the mixed frame
+        :math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
+        active velocity representation is either inertial-fixed or mixed,
+        and :math:`[C] = [B]` if the active velocity representation is body-fixed.
+    """
+
+    ν = data.generalized_velocity()
+    G_J = average_centroidal_velocity_jacobian(model=model, data=data)
+
+    return G_J @ ν
+
+
+@jax.jit
+def average_centroidal_velocity_jacobian(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Matrix:
+    r"""
+    Compute the Jacobian of the average centroidal velocity of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The Jacobian of the average centroidal velocity of the model.
+
+    Note:
+        The frame corresponding to the output representation of this Jacobian is either
+        :math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
+        or :math:`G[B]`, if the active velocity representation is body-fixed.
+    """
+
+    G_J = centroidal_momentum_jacobian(model=model, data=data)
+    G_Mbb = locked_centroidal_spatial_inertia(model=model, data=data)
+
+    return jnp.linalg.inv(G_Mbb) @ G_J

jaxsim/api/common.py

@@ -0,0 +1,216 @@
+import abc
+import contextlib
+import dataclasses
+import enum
+import functools
+from typing import ContextManager
+
+import jax
+import jax.numpy as jnp
+import jax_dataclasses
+import jaxlie
+from jax_dataclasses import Static
+
+import jaxsim.typing as jtp
+from jaxsim.utils import JaxsimDataclass, Mutability
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+
+@enum.unique
+class VelRepr(enum.IntEnum):
+    """
+    Enumeration of all supported 6D velocity representations.
+    """
+
+    Body = enum.auto()
+    Mixed = enum.auto()
+    Inertial = enum.auto()
+
+
+@jax_dataclasses.pytree_dataclass
+class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
+    """
+    Base class for model data structures with velocity representation.
+    """
+
+    velocity_representation: Static[VelRepr] = dataclasses.field(
+        default=VelRepr.Inertial, kw_only=True
+    )
+
+    @contextlib.contextmanager
+    def switch_velocity_representation(
+        self, velocity_representation: VelRepr
+    ) -> ContextManager[Self]:
+        """
+        Context manager to temporarily switch the velocity representation.
+
+        Args:
+            velocity_representation: The new velocity representation.
+
+        Yields:
+            The same object with the new velocity representation.
+        """
+
+        original_representation = self.velocity_representation
+
+        try:
+
+            # First, we replace the velocity representation
+            with self.mutable_context(
+                mutability=Mutability.MUTABLE_NO_VALIDATION,
+                restore_after_exception=True,
+            ):
+                self.velocity_representation = velocity_representation
+
+            # Then, we yield the data with changed representation.
+            # We run this in a mutable context with restoration so that any exception
+            # occurring, we restore the original object in case it was modified.
+            with self.mutable_context(
+                mutability=self.mutability(), restore_after_exception=True
+            ):
+                yield self
+
+        finally:
+            with self.mutable_context(
+                mutability=Mutability.MUTABLE_NO_VALIDATION,
+                restore_after_exception=True,
+            ):
+                self.velocity_representation = original_representation
+
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=["other_representation", "is_force"])
+    def inertial_to_other_representation(
+        array: jtp.Array,
+        other_representation: VelRepr,
+        transform: jtp.Matrix,
+        is_force: bool = False,
+    ) -> jtp.Array:
+        r"""
+        Convert a 6D quantity from inertial-fixed to another representation.
+
+        Args:
+            array: The 6D quantity to convert.
+            other_representation: The representation to convert to.
+            transform:
+                The `math:W \mathbf{H}_O` transform, where `math:O` is the
+                reference frame of the other representation.
+            is_force: Whether the quantity is a 6D force or a 6D velocity.
+
+        Returns:
+            The 6D quantity in the other representation.
+        """
+
+        W_array = array.squeeze()
+        W_H_O = transform.squeeze()
+
+        if W_array.size != 6:
+            raise ValueError(W_array.size, 6)
+
+        if W_H_O.shape != (4, 4):
+            raise ValueError(W_H_O.shape, (4, 4))
+
+        match other_representation:
+
+            case VelRepr.Inertial:
+                return W_array
+
+            case VelRepr.Body:
+
+                if not is_force:
+                    O_Xv_W = jaxlie.SE3.from_matrix(W_H_O).inverse().adjoint()
+                    O_array = O_Xv_W @ W_array
+
+                else:
+                    O_Xf_W = jaxlie.SE3.from_matrix(W_H_O).adjoint().T
+                    O_array = O_Xf_W @ W_array
+
+                return O_array
+
+            case VelRepr.Mixed:
+                W_p_O = W_H_O[0:3, 3]
+                W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
+
+                if not is_force:
+                    OW_Xv_W = jaxlie.SE3.from_matrix(W_H_OW).inverse().adjoint()
+                    OW_array = OW_Xv_W @ W_array
+
+                else:
+                    OW_Xf_W = jaxlie.SE3.from_matrix(W_H_OW).adjoint().transpose()
+                    OW_array = OW_Xf_W @ W_array
+
+                return OW_array
+
+            case _:
+                raise ValueError(other_representation)
+
+    @staticmethod
+    @functools.partial(jax.jit, static_argnames=["other_representation", "is_force"])
+    def other_representation_to_inertial(
+        array: jtp.Array,
+        other_representation: VelRepr,
+        transform: jtp.Matrix,
+        is_force: bool = False,
+    ) -> jtp.Array:
+        r"""
+        Convert a 6D quantity from another representation to inertial-fixed.
+
+        Args:
+            array: The 6D quantity to convert.
+            other_representation: The representation to convert from.
+            transform:
+                The `math:W \mathbf{H}_O` transform, where `math:O` is the
+                reference frame of the other representation.
+            is_force: Whether the quantity is a 6D force or a 6D velocity.
+
+        Returns:
+            The 6D quantity in the inertial-fixed representation.
+        """
+
+        W_array = array.squeeze()
+        W_H_O = transform.squeeze()
+
+        if W_array.size != 6:
+            raise ValueError(W_array.size, 6)
+
+        if W_H_O.shape != (4, 4):
+            raise ValueError(W_H_O.shape, (4, 4))
+
+        match other_representation:
+            case VelRepr.Inertial:
+                W_array = array
+                return W_array
+
+            case VelRepr.Body:
+                O_array = array
+
+                if not is_force:
+                    W_Xv_O: jtp.Array = jaxlie.SE3.from_matrix(W_H_O).adjoint()
+                    W_array = W_Xv_O @ O_array
+
+                else:
+                    W_Xf_O = jaxlie.SE3.from_matrix(W_H_O).inverse().adjoint().T
+                    W_array = W_Xf_O @ O_array
+
+                return W_array
+
+            case VelRepr.Mixed:
+                BW_array = array
+                W_p_O = W_H_O[0:3, 3]
+                W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
+
+                if not is_force:
+                    W_Xv_BW: jtp.Array = jaxlie.SE3.from_matrix(W_H_OW).adjoint()
+                    W_array = W_Xv_BW @ BW_array
+
+                else:
+                    W_Xf_BW = jaxlie.SE3.from_matrix(W_H_OW).inverse().adjoint().T
+                    W_array = W_Xf_BW @ BW_array
+
+                return W_array
+
+            case _:
+                raise ValueError(other_representation)