jaxsim 0.5.1.dev4__tar.gz → 0.5.1.dev15__tar.gz

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev4
+Version: 0.5.1.dev15
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/examples/jaxsim_as_multibody_dynamics_library.ipynb

@@ -320,12 +320,18 @@
     "\n",
     "- **`jaxsim.api.model`**: vectorized functions operating on the whole model.\n",
     "- **`jaxsim.api.link`**: functions operating on individual links.\n",
+    "- **`jaxsim.api.frame`**: functions operating on individual frames. \n",
     "\n",
-    "Due to JAX limitations on vectorizable data types, many APIs operate on indices instead of names. Since using indices can be error prone, JaxSim provides conversion functions:\n",
+    "Due to JAX limitations on vectorizable data types, many APIs operate on indices instead of names. Since using indices can be error prone, JaxSim provides conversion functions for both links:\n",
     "\n",
     "- **jaxsim.api.link.names_to_idxs()**\n",
     "- **jaxsim.api.link.idxs_to_names()**\n",
     "\n",
+    "and frames: \n",
+    "\n",
+    "- **jaxsim.api.frame.names_to_idxs()**\n",
+    "- **jaxsim.api.frame.idxs_to_names()**\n",
+    "\n",
     "We recommend using names whenever possible to avoid hard-to-trace errors.\n"
    ]
   },
@@ -354,7 +360,7 @@
    },
    "outputs": [],
    "source": [
-    "# @title Pose\n",
+    "# @title Link Pose\n",
     "\n",
     "# Compute its pose w.r.t. the world frame through forward kinematics.\n",
     "W_H_L = js.link.transform(model=model, data=data, link_index=link_index)\n",
@@ -380,7 +386,7 @@
    },
    "outputs": [],
    "source": [
-    "# @title 6D Velocity\n",
+    "# @title Link 6D Velocity\n",
     "\n",
     "# JaxSim allows to select the so-called representation of the frame velocity.\n",
     "L_v_WL = js.link.velocity(model=model, data=data, link_index=link_index, output_vel_repr=VelRepr.Body)\n",
@@ -413,10 +419,87 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {
     "id": "SSoziCShtwZ9"
    },
+   "outputs": [],
+   "source": [
+    "# Find the index of a frame.\n",
+    "frame_name = \"l_foot_front\"\n",
+    "frame_index = js.frame.name_to_idx(model=model, frame_name=frame_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "fVp_xP_1twZ9",
+    "outputId": "cfaa0569-d768-4708-c98c-a5867c056d04"
+   },
+   "outputs": [],
+   "source": [
+    "# @title Frame Pose\n",
+    "\n",
+    "# Compute its pose w.r.t. the world frame through forward kinematics.\n",
+    "W_H_F = js.frame.transform(model=model, data=data, frame_index=frame_index)\n",
+    "\n",
+    "print(f\"Transform of '{frame_name}': shape={W_H_F.shape}\\n{W_H_F}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "QqaqxneEFYiW"
+   },
+   "outputs": [],
+   "source": [
+    "# @title Frame 6D Velocity\n",
+    "\n",
+    "# JaxSim allows to select the so-called representation of the frame velocity.\n",
+    "F_v_WF = js.frame.velocity(model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Body)\n",
+    "FW_v_WF = js.frame.velocity(model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Mixed)\n",
+    "W_v_WF = js.frame.velocity(model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Inertial)\n",
+    "\n",
+    "print(f\"Body-fixed velocity      F_v_WF={F_v_WF}\")\n",
+    "print(f\"Mixed velocity:         FW_v_WF={FW_v_WF}\")\n",
+    "print(f\"Inertial-fixed velocity: W_v_WF={W_v_WF}\")\n",
+    "\n",
+    "# These can also be computed passing through the frame free-floating Jacobian.\n",
+    "# This type of Jacobian has a input velocity representation that corresponds\n",
+    "# the velocity representation of ν, and an output velocity representation that\n",
+    "# corresponds to the velocity representation of the desired 6D velocity.\n",
+    "\n",
+    "# You can use the following context manager to easily switch between representations.\n",
+    "with data.switch_velocity_representation(VelRepr.Body):\n",
+    "\n",
+    "    # Body-fixed generalized velocity.\n",
+    "    B_ν = data.generalized_velocity()\n",
+    "\n",
+    "    # Free-floating Jacobian accepting a body-fixed generalized velocity and\n",
+    "    # returning an inertial-fixed link velocity.\n",
+    "    W_J_WF_B = js.frame.jacobian(\n",
+    "        model=model, data=data, frame_index=frame_index, output_vel_repr=VelRepr.Inertial\n",
+    "    )\n",
+    "\n",
+    "# Now the following relation should hold.\n",
+    "assert jnp.allclose(W_v_WF, W_J_WF_B @ B_ν)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "d_vp6D74GoVZ",
+    "outputId": "798b9283-792e-4339-b56c-df2595fac974"
+   },
    "source": [
     "## Robot Dynamics\n",
     "\n",
@@ -460,11 +543,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "fVp_xP_1twZ9",
-    "outputId": "cfaa0569-d768-4708-c98c-a5867c056d04"
+    "id": "oOKJOVfsH4Ki"
    },
    "outputs": [],
    "source": [
@@ -509,7 +588,11 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "QqaqxneEFYiW"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "FlNo8dNWKKtu",
+    "outputId": "313e939b-f88f-4407-c9ee-b5b3b7443061"
    },
    "source": [
     "### Forward Dynamics\n",
@@ -530,11 +613,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "d_vp6D74GoVZ",
-    "outputId": "798b9283-792e-4339-b56c-df2595fac974"
+    "id": "LXARuRu1Ly1K"
    },
    "outputs": [],
    "source": [
@@ -554,7 +633,11 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "oOKJOVfsH4Ki"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "g5GOYXDnLySU",
+    "outputId": "ad4ce77d-d06f-473a-9c32-040680d76aa5"
    },
    "source": [
     "### Inverse Dynamics\n",
@@ -575,11 +658,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "FlNo8dNWKKtu",
-    "outputId": "313e939b-f88f-4407-c9ee-b5b3b7443061"
+    "id": "UTae5MjhaP2H"
    },
    "outputs": [],
    "source": [
@@ -604,7 +683,11 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "LXARuRu1Ly1K"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "gYZ1jK1Neg1H",
+    "outputId": "0de79770-1e18-4027-bb47-5713bc1b4a72"
    },
    "source": [
     "### Centroidal Dynamics\n",
@@ -656,11 +739,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "g5GOYXDnLySU",
-    "outputId": "ad4ce77d-d06f-473a-9c32-040680d76aa5"
+    "id": "rrSfxp8lh9YZ"
    },
    "outputs": [],
    "source": [
@@ -716,7 +795,11 @@
   {
    "cell_type": "markdown",
    "metadata": {
-    "id": "UTae5MjhaP2H"
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "Ot6HePB_twaE",
+    "outputId": "02a6abae-257e-45ee-e9de-6a607cdbeb9a"
    },
    "source": [
     "## Contact Frames\n",
@@ -746,11 +829,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "gYZ1jK1Neg1H",
-    "outputId": "0de79770-1e18-4027-bb47-5713bc1b4a72"
+    "id": "LITRC3STliKR"
    },
    "outputs": [],
    "source": [

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/pyproject.toml

@@ -268,3 +268,6 @@ cuda-version = "12.*"
 # Pinning a specific version awaiting the following to get addressed:
 # https://github.com/conda-forge/jaxlib-feedstock/issues/285
 jaxlib = { version = "<0.4.31", build = "*cuda*" }
+
+[tool.pixi.feature.gpu.tasks]
+gpu-tests = { cmd = "pytest --gpu-only", depends_on = ["pipcheck"] }

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev4'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev4')
+__version__ = version = '0.5.1.dev15'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev15')

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/src/jaxsim/api/frame.py

@@ -58,7 +58,7 @@ def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> jtp.Int:
         The index of the frame.
     """
 
-    if frame_name not in model.kin_dyn_parameters.frame_parameters.name:
+    if frame_name not in model.frame_names():
         raise ValueError(f"Frame '{frame_name}' not found in the model.")
 
     return (
@@ -180,6 +180,56 @@ def transform(
     return W_H_L @ L_H_F
 
 
+@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+def velocity(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    frame_index: jtp.IntLike,
+    output_vel_repr: VelRepr | None = None,
+) -> jtp.Vector:
+    """
+    Compute the 6D velocity of the frame.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        frame_index: The index of the frame.
+        output_vel_repr:
+            The output velocity representation of the frame velocity.
+
+    Returns:
+        The 6D velocity of the frame in the specified velocity representation.
+    """
+    n_l = model.number_of_links()
+    n_f = model.number_of_frames()
+
+    exceptions.raise_value_error_if(
+        condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),
+        msg="Invalid frame index '{idx}'",
+        idx=frame_index,
+    )
+
+    output_vel_repr = (
+        output_vel_repr if output_vel_repr is not None else data.velocity_representation
+    )
+
+    # Get the frame jacobian having I as input representation (taken from data)
+    # and O as output representation, specified by the user (or taken from data).
+    O_J_WF_I = jacobian(
+        model=model,
+        data=data,
+        frame_index=frame_index,
+        output_vel_repr=output_vel_repr,
+    )
+
+    # Get the generalized velocity in the input velocity representation.
+    I_ν = data.generalized_velocity()
+
+    # Compute the frame velocity in the output velocity representation.
+    return O_J_WF_I @ I_ν
+
+
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian(
     model: js.model.JaxSimModel,
@@ -207,7 +257,7 @@ def jacobian(
     """
 
     n_l = model.number_of_links()
-    n_f = len(model.frame_names())
+    n_f = model.number_of_frames()
 
     exceptions.raise_value_error_if(
         condition=jnp.array([frame_index < n_l, frame_index >= n_l + n_f]).any(),

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/src/jaxsim/api/kin_dyn_parameters.py

@@ -268,6 +268,16 @@ class KynDynParameters(JaxsimDataclass):
 
         return len(self.joint_model.joint_names) - 1
 
+    def number_of_frames(self) -> int:
+        """
+        Return the number of frames of the model.
+
+        Returns:
+            The number of frames of the model.
+        """
+
+        return len(self.frame_parameters.name)
+
     def support_body_array(self, link_index: jtp.IntLike) -> jtp.Vector:
         r"""
         Return the support parent array :math:`\kappa(i)` of a link.

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/src/jaxsim/api/model.py

@@ -327,6 +327,17 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.kin_dyn_parameters.number_of_joints()
 
+    def number_of_frames(self) -> int:
+        """
+        Return the number of frames in the model.
+
+        Returns:
+            The number of frames in the model.
+
+        """
+
+        return self.kin_dyn_parameters.number_of_frames()
+
     # =================
     # Base link methods
     # =================

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev4
+Version: 0.5.1.dev15
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/tests/conftest.py

@@ -1,5 +1,6 @@
 import os
 import pathlib
+import subprocess
 
 import jax
 import pytest
@@ -11,7 +12,45 @@ import jaxsim
 import jaxsim.api as js
 
 
-def pytest_configure() -> None:
+def pytest_addoption(parser):
+    parser.addoption(
+        "--gpu-only",
+        action="store_true",
+        default=False,
+        help="Run tests only if GPU is available and utilized",
+    )
+
+
+def check_gpu_usage():
+    # Set environment variable to prioritize GPU.
+    os.environ["JAX_PLATFORM_NAME"] = "gpu"
+
+    # Run a simple JAX operation
+    x = jax.device_put(jax.numpy.ones((512, 512)))
+    y = jax.device_put(jax.numpy.ones((512, 512)))
+    _ = jax.numpy.dot(x, y).block_until_ready()
+
+    # Check GPU memory usage with nvidia-smi.
+    result = subprocess.run(
+        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader"],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        pytest.exit(
+            "Failed to query GPU usage. Ensure nvidia-smi is installed and accessible."
+        )
+
+    gpu_memory_usage = [
+        int(line.strip().split()[0]) for line in result.stdout.splitlines()
+    ]
+    if all(usage == 0 for usage in gpu_memory_usage):
+        pytest.exit(
+            "GPU is available but not utilized during computations. Check your JAX installation."
+        )
+
+
+def pytest_configure(config) -> None:
     """Pytest configuration hook."""
 
     # This is a global variable that is updated by the `prng_key` fixture.
@@ -19,6 +58,15 @@ def pytest_configure() -> None:
         seed=int(os.environ.get("JAXSIM_TEST_SEED", 0))
     )
 
+    # Check if GPU is available and utilized.
+    if config.getoption("--gpu-only"):
+        devices = jax.devices()
+        if not any(device.platform == "gpu" for device in devices):
+            pytest.exit("No GPU devices found. Check your JAX installation.")
+
+        # Ensure GPU is being used during computation
+        check_gpu_usage()
+
 
 # ================
 # Generic fixtures

{jaxsim-0.5.1.dev4 → jaxsim-0.5.1.dev15}/tests/test_api_frame.py

@@ -171,6 +171,12 @@ def test_frame_jacobians(
         J_WL_idt = kin_dyn.jacobian_frame(frame_name=frame_name)
         assert J_WL_js == pytest.approx(J_WL_idt, abs=1e-9)
 
+    for frame_name, frame_index in zip(frame_names, frame_indices, strict=True):
+
+        v_WF_idt = kin_dyn.frame_velocity(frame_name=frame_name)
+        v_WF_js = js.frame.velocity(model=model, data=data, frame_index=frame_index)
+        assert v_WF_js == pytest.approx(v_WF_idt), frame_name
+
 
 def test_frame_jacobian_derivative(
     jaxsim_models_types: js.model.JaxSimModel,