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

jaxsim/__init__.py

@@ -8,16 +8,37 @@ def _jnp_options() -> None:
 
     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")
+    # Check if running on TPU.
+    is_tpu = jax.devices()[0].platform == "tpu"
+
+    # Check if running on Metal.
+    is_metal = jax.devices()[0].platform == "METAL"
+
+    # Enable by default 64-bit precision to get accurate physics.
+    # Users can enforce 32-bit precision by setting the following variable to 0.
+    use_x64 = os.environ.get("JAX_ENABLE_X64", "1") != "0"
+
+    # Notify the user if unsupported 64-bit precision was enforced on TPU.
+    if (is_tpu or is_metal) and use_x64:
+        msg = f"64-bit precision is not allowed on {jax.devices()[0].platform.upper}. Enforcing 32bit precision."
+        logging.warning(msg)
+        use_x64 = False
+
+        if is_metal:
+            logging.warning(
+                "JAX Metal backend is experimental. Some functionalities may not be available."
+            )
+
+    # Enable 64-bit precision in JAX.
+    if use_x64:
+        logging.info("Enabling JAX to use 64-bit precision")
         jax.config.update("jax_enable_x64", True)
 
-        import jax.numpy as jnp
-        import numpy as np
-
-        if jnp.empty(0, dtype=float).dtype != jnp.empty(0, dtype=np.float64).dtype:
-            logging.warning("Failed to enable 64bit precision in JAX")
+    # Warn about experimental usage of 32-bit precision.
+    else:
+        logging.warning(
+            "Using 32-bit precision in JaxSim is still experimental, please avoid to use variable step integrators."
+        )
 
 
 def _np_options() -> None:
@@ -27,41 +48,71 @@ def _np_options() -> None:
 
 
 def _is_editable() -> bool:
+
     import importlib.util
     import pathlib
     import site
 
-    # Get the ModuleSpec of jaxsim
+    # Get the ModuleSpec of jaxsim.
     jaxsim_spec = importlib.util.find_spec(name="jaxsim")
 
     # This can be None. If it's None, assume non-editable installation.
     if jaxsim_spec.origin is None:
         return False
 
-    # Get the folder containing the jaxsim package
+    # Get the folder containing the jaxsim package.
     jaxsim_package_dir = str(pathlib.Path(jaxsim_spec.origin).parent.parent)
 
-    # The installation is editable if the package dir is not in any {site|dist}-packages
+    # The installation is editable if the package dir is not in any {site|dist}-packages.
     return jaxsim_package_dir not in site.getsitepackages()
 
 
-# Initialize the logging verbosity
-if _is_editable():
-    logging.configure(level=logging.LoggingLevel.DEBUG)
-else:
-    logging.configure(level=logging.LoggingLevel.WARNING)
+def _get_default_logging_level(env_var: str) -> logging.LoggingLevel:
+    """
+    Get the default logging level.
+
+    Args:
+        env_var: The environment variable to check.
+
+    Returns:
+        The logging level to set.
+    """
+
+    import os
+
+    # Define the default logging level depending on the installation mode.
+    default_logging_level = (
+        logging.LoggingLevel.DEBUG
+        if _is_editable()  # noqa: F821
+        else logging.LoggingLevel.WARNING
+    )
+
+    # Allow to override the default logging level with an environment variable.
+    try:
+        return logging.LoggingLevel[
+            os.environ.get(env_var, default_logging_level.name).upper()
+        ]
+
+    except KeyError as exc:
+        msg = f"Invalid logging level defined in {env_var}='{os.environ[env_var]}'"
+        raise RuntimeError(msg) from exc
+
+
+# Configure the logger with the default logging level.
+logging.configure(level=_get_default_logging_level(env_var="JAXSIM_LOGGING_LEVEL"))
+
 
-# Configure JAX
+# Configure JAX.
 _jnp_options()
 
-# Initialize the numpy print options
+# Initialize the numpy print options.
 _np_options()
 
 del _jnp_options
 del _np_options
+del _get_default_logging_level
 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.2.dev191'
-__version_tuple__ = version_tuple = (0, 2, 'dev191')
+__version__ = version = '0.6.1.dev2'
+__version_tuple__ = version_tuple = (0, 6, 1, 'dev2')

jaxsim/api/__init__.py

@@ -1 +1,13 @@
-from . import contact, data, joint, link, model, ode
+from . import common  # isort:skip
+from . import model, data  # isort:skip
+from . import (
+    com,
+    contact,
+    frame,
+    joint,
+    kin_dyn_parameters,
+    link,
+    ode,
+    ode_data,
+    references,
+)

jaxsim/api/com.py

@@ -0,0 +1,423 @@
+import jax
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.math
+import jaxsim.typing as jtp
+
+from .common import VelRepr
+
+
+@jax.jit
+@js.common.named_scope
+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 = jaxsim.math.Transform.inverse(transform=W_H_B)
+
+    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
+@js.common.named_scope
+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
+@js.common.named_scope
+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
+@js.common.named_scope
+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)  # noqa: F841
+        case VelRepr.Body:
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)  # noqa: F841
+        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
+@js.common.named_scope
+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)  # noqa: F841
+        case VelRepr.Body:
+            W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)  # noqa: F841
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    B_H_G = jaxsim.math.Transform.inverse(W_H_B) @ W_H_G
+
+    B_Xv_G = jaxsim.math.Adjoint.from_transform(transform=B_H_G)
+    G_Xf_B = B_Xv_G.transpose()
+
+    return G_Xf_B @ B_Mbb_B @ B_Xv_G
+
+
+@jax.jit
+@js.common.named_scope
+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
+@js.common.named_scope
+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
+
+
+@jax.jit
+@js.common.named_scope
+def bias_acceleration(
+    model: js.model.JaxSimModel, data: js.data.JaxSimModelData
+) -> jtp.Vector:
+    r"""
+    Compute the bias linear acceleration of the center of mass.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The bias linear acceleration of the center of mass in the active representation.
+
+    Note:
+        The bias acceleration 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.
+    """
+
+    # Compute the pose of all links with forward kinematics.
+    W_H_L = js.model.forward_kinematics(model=model, data=data)
+
+    # Compute the bias acceleration of all links by zeroing the generalized velocity
+    # in the active representation.
+    v̇_bias_WL = js.model.link_bias_accelerations(model=model, data=data)
+
+    def other_representation_to_body(
+        C_v̇_WL: jtp.Vector, C_v_WC: jtp.Vector, L_H_C: jtp.Matrix, L_v_LC: jtp.Vector
+    ) -> jtp.Vector:
+        """
+        Convert the body-fixed representation of the link bias acceleration
+        C_v̇_WL expressed in a generic frame C to the body-fixed representation L_v̇_WL.
+        """
+
+        L_X_C = jaxsim.math.Adjoint.from_transform(transform=L_H_C)
+        C_X_L = jaxsim.math.Adjoint.inverse(L_X_C)
+
+        L_v̇_WL = L_X_C @ (C_v̇_WL + jaxsim.math.Cross.vx(C_X_L @ L_v_LC) @ C_v_WC)
+        return L_v̇_WL
+
+    # We need here to get the body-fixed bias acceleration of the links.
+    # Since it's computed in the active representation, we need to convert it to body.
+    match data.velocity_representation:
+
+        case VelRepr.Body:
+            L_a_bias_WL = v̇_bias_WL
+
+        case VelRepr.Inertial:
+
+            C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL  # noqa: F841
+            C_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
+
+            L_H_C = L_H_W = jax.vmap(  # noqa: F841
+                lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
+            )(W_H_L)
+
+            L_v_LC = L_v_LW = jax.vmap(  # noqa: F841
+                lambda i: -js.link.velocity(
+                    model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
+                )
+            )(jnp.arange(model.number_of_links()))
+
+            L_a_bias_WL = jax.vmap(
+                lambda i: other_representation_to_body(
+                    C_v̇_WL=C_v̇_WL[i],
+                    C_v_WC=C_v_WC,
+                    L_H_C=L_H_C[i],
+                    L_v_LC=L_v_LC[i],
+                )
+            )(jnp.arange(model.number_of_links()))
+
+        case VelRepr.Mixed:
+
+            C_v̇_WL = LW_v̇_bias_WL = v̇_bias_WL  # noqa: F841
+
+            C_v_WC = LW_v_W_LW = jax.vmap(  # noqa: F841
+                lambda i: js.link.velocity(
+                    model=model, data=data, link_index=i, output_vel_repr=VelRepr.Mixed
+                )
+                .at[3:6]
+                .set(jnp.zeros(3))
+            )(jnp.arange(model.number_of_links()))
+
+            L_H_C = L_H_LW = jax.vmap(  # noqa: F841
+                lambda W_H_L: jaxsim.math.Transform.inverse(
+                    W_H_L.at[0:3, 3].set(jnp.zeros(3))
+                )
+            )(W_H_L)
+
+            L_v_LC = L_v_L_LW = jax.vmap(  # noqa: F841
+                lambda i: -js.link.velocity(
+                    model=model, data=data, link_index=i, output_vel_repr=VelRepr.Body
+                )
+                .at[0:3]
+                .set(jnp.zeros(3))
+            )(jnp.arange(model.number_of_links()))
+
+            L_a_bias_WL = jax.vmap(
+                lambda i: other_representation_to_body(
+                    C_v̇_WL=C_v̇_WL[i],
+                    C_v_WC=C_v_WC[i],
+                    L_H_C=L_H_C[i],
+                    L_v_LC=L_v_LC[i],
+                )
+            )(jnp.arange(model.number_of_links()))
+
+        case _:
+            raise ValueError(data.velocity_representation)
+
+    # Compute the bias of the 6D momentum derivative.
+    def bias_momentum_derivative_term(
+        link_index: jtp.Int, L_a_bias_WL: jtp.Vector
+    ) -> jtp.Vector:
+
+        # Get the body-fixed 6D inertia matrix.
+        L_M_L = js.link.spatial_inertia(model=model, link_index=link_index)
+
+        # Compute the body-fixed 6D velocity.
+        L_v_WL = js.link.velocity(
+            model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Body
+        )
+
+        # Compute the world-to-link transformations for 6D forces.
+        W_Xf_L = jaxsim.math.Adjoint.from_transform(
+            transform=W_H_L[link_index], inverse=True
+        ).T
+
+        # Compute the contribution of the link to the bias acceleration of the CoM.
+        W_ḣ_bias_link_contribution = W_Xf_L @ (
+            L_M_L @ L_a_bias_WL + jaxsim.math.Cross.vx_star(L_v_WL) @ L_M_L @ L_v_WL
+        )
+
+        return W_ḣ_bias_link_contribution
+
+    # Sum the contributions of all links to the bias acceleration of the CoM.
+    W_ḣ_bias = jax.vmap(bias_momentum_derivative_term)(
+        jnp.arange(model.number_of_links()), L_a_bias_WL
+    ).sum(axis=0)
+
+    # Compute the total mass of the model.
+    m = js.model.total_mass(model=model)
+
+    # Compute the position of the CoM.
+    W_p_CoM = com_position(model=model, data=data)
+
+    match data.velocity_representation:
+
+        # G := G[W] = (W_p_CoM, [W])
+        case VelRepr.Inertial | VelRepr.Mixed:
+
+            W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
+            GW_Xf_W = jaxsim.math.Adjoint.from_transform(W_H_GW).T
+
+            GW_ḣ_bias = GW_Xf_W @ W_ḣ_bias
+            GW_v̇l_com_bias = GW_ḣ_bias[0:3] / m
+
+            return GW_v̇l_com_bias
+
+        # G := G[B] = (W_p_CoM, [B])
+        case VelRepr.Body:
+
+            GB_Xf_W = jaxsim.math.Adjoint.from_transform(
+                transform=data.base_transform().at[0:3].set(W_p_CoM)
+            ).T
+
+            GB_ḣ_bias = GB_Xf_W @ W_ḣ_bias
+            GB_v̇l_com_bias = GB_ḣ_bias[0:3] / m
+
+            return GB_v̇l_com_bias
+
+        case _:
+            raise ValueError(data.velocity_representation)