jaxsim 0.2.1.dev56__tar.gz → 0.2.1.dev69__tar.gz

{jaxsim-0.2.1.dev56 → jaxsim-0.2.1.dev69}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.2.1.dev56
+Version: 0.2.1.dev69
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.2.1.dev56 → jaxsim-0.2.1.dev69}/examples/PD_controller.ipynb

@@ -6,7 +6,11 @@
    "source": [
     "# `JAXsim` Showcase: PD Controller\n",
     "\n",
-    "First, we install the necessary packages and import them."
+    "First, we install the necessary packages and import them.\n",
+    "\n",
+    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/PD_controller.ipynb\">\n",
+    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
+    "</a>"
    ]
   },
   {
@@ -19,6 +23,28 @@
     "from IPython.display import clear_output, HTML, display\n",
     "import sys\n",
     "\n",
+    "IS_COLAB = \"google.colab\" in sys.modules\n",
+    "\n",
+    "# Install JAX and Gazebo SDF\n",
+    "if IS_COLAB:\n",
+    "    !{sys.executable} -m pip install -qU jaxsim[viz]\n",
+    "    !apt install -qq lsb-release wget gnupg\n",
+    "    !wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg\n",
+    "    !echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null\n",
+    "    !apt -qq update\n",
+    "    !apt install -qq --no-install-recommends libsdformat13 gz-tools2\n",
+    "\n",
+    "    # Optional dependency for visualization\n",
+    "    !{sys.executable} -m pip install pyvirtualdisplay\n",
+    "    !apt install xvfb\n",
+    "\n",
+    "    # Start Xvfb and set \"DISPLAY\" environment variable\n",
+    "    import pyvirtualdisplay\n",
+    "    d = pyvirtualdisplay.Display()\n",
+    "    d.start()\n",
+    "\n",
+    "    clear_output()\n",
+    "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "from jaxsim import logging\n",
@@ -124,29 +150,50 @@
    "source": [
     "# @title Set up MuJoCo renderer\n",
     "\n",
-    "from jaxsim.mujoco.visualizer import MujocoVisualizer\n",
     "from jaxsim.mujoco import RodModelToMjcf, MujocoModelHelper, MujocoVideoRecorder\n",
     "from jaxsim.mujoco.loaders import UrdfToMjcf\n",
     "\n",
+    "import os\n",
+    "import subprocess\n",
+    "\n",
+    "\n",
+    "if IS_COLAB:\n",
+    "    if subprocess.run(\"ffmpeg -version\", shell=True).returncode:\n",
+    "        !command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
+    "        clear_output()\n",
+    "\n",
+    "    if subprocess.run(\"nvidia-smi\").returncode:\n",
+    "        raise RuntimeError(\n",
+    "            \"Cannot communicate with GPU. \"\n",
+    "            \"Make sure you are using a GPU Colab runtime. \"\n",
+    "            \"Go to the Runtime menu and select Choose runtime type.\"\n",
+    "        )\n",
+    "\n",
+    "    # Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
+    "    # This is usually installed as part of an Nvidia driver package, but the Colab\n",
+    "    # kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
+    "    # (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
+    "    NVIDIA_ICD_CONFIG_PATH = \"/usr/share/glvnd/egl_vendor.d/10_nvidia.json\"\n",
+    "    if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
+    "        with open(NVIDIA_ICD_CONFIG_PATH, \"w\") as f:\n",
+    "            f.write(\n",
+    "                \"\"\"{\n",
+    "      \"file_format_version\" : \"1.0.0\",\n",
+    "      \"ICD\" : {\n",
+    "         \"library_path\" : \"libEGL_nvidia.so.0\"\n",
+    "     }\n",
+    "  }\n",
+    "  \"\"\"\n",
+    "            )\n",
     "\n",
     "%env MUJOCO_GL=egl\n",
     "\n",
-    "try:\n",
-    "    import mujoco\n",
-    "except Exception as e:\n",
-    "    raise e from RuntimeError(\n",
-    "        \"Something went wrong during installation. Check the shell output above \"\n",
-    "        \"for more information.\\n\"\n",
-    "        \"If using a hosted Colab runtime, make sure you enable GPU acceleration \"\n",
-    "        'by going to the Runtime menu and selecting \"Choose runtime type\".'\n",
-    "    )\n",
-    "\n",
     "camera = {\n",
-    "    \"name\":\"cartpole_camera\",\n",
-    "    \"mode\":\"fixed\",\n",
-    "    \"pos\":\"3.954 3.533 2.343\",\n",
-    "    \"xyaxes\":\"-0.594 0.804 -0.000 -0.163 -0.120 0.979\",\n",
-    "    \"fovy\":\"60\",\n",
+    "    \"name\": \"cartpole_camera\",\n",
+    "    \"mode\": \"fixed\",\n",
+    "    \"pos\": \"3.954 3.533 2.343\",\n",
+    "    \"xyaxes\": \"-0.594 0.804 -0.000 -0.163 -0.120 0.979\",\n",
+    "    \"fovy\": \"60\",\n",
     "}\n",
     "\n",
     "mjcf_string, assets = UrdfToMjcf.convert(urdf=model.built_from, cameras=camera)\n",

{jaxsim-0.2.1.dev56 → jaxsim-0.2.1.dev69}/examples/Parallel_computing.ipynb

@@ -25,7 +25,16 @@
     "\n",
     "from IPython.display import HTML, clear_output, display\n",
     "\n",
+    "IS_COLAB = \"google.colab\" in sys.modules\n",
+    "\n",
     "# Install JAX and Gazebo\n",
+    "if IS_COLAB:\n",
+    "    !{sys.executable} -m pip install -qU jaxsim\n",
+    "    !apt install -qq lsb-release wget gnupg\n",
+    "    !wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg\n",
+    "    !echo \"deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main\" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null\n",
+    "    !apt -qq update\n",
+    "    !apt install -qq --no-install-recommends libsdformat13 gz-tools2\n",
     "\n",
     "# Set environment variable to avoid GPU out of memory errors\n",
     "%env XLA_PYTHON_CLIENT_MEM_PREALLOCATE=false\n",

jaxsim-0.2.1.dev69/examples/README.md

@@ -0,0 +1,41 @@
+# JAXsim Notebook Examples
+
+This folder includes a Jupyter Notebook demonstrating the practical usage of JAXsim for system simulations.
+
+### Examples
+
+- [PD_controller](./PD_controller.ipynb) <a target="_blank" href="https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/PD_controller.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a> - A simple example demonstrating the use of JAXsim to simulate a PD controller with gravity compensation for a 2-DOF cartpole.
+- [Parallel_computing](./Parallel_computing.ipynb) <a target="_blank" href="https://colab.research.google.com/github/ami-iit/jaxsim/blob/main/examples/Parallel_computing.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a> - An example demonstrating how to simulate vectorized models in parallel using JAXsim.
+
+> [!TIP]
+> Stay tuned for more examples!
+
+## Running the Examples
+
+To execute these examples utilizing JAXsim with hardware acceleration, there are a couple of options available:
+
+### Option 1: Google Colab (Recommended)
+
+The simplest way to run the examples is by accessing the provided Google Colab notebook link mentioned above. This will enable you to execute the examples in a hosted environment.
+
+### Option 2: Local Execution with `pixi`
+
+For local execution, follow these steps:
+
+1. **Install `pixi`:**
+
+As per the [official documentation](https://pixi.sh/#installation):
+
+```bash
+curl -fsSL https://pixi.sh/install.sh | bash
+```
+
+2. **Run the Example Notebook:**
+
+Use `pixi run examples` from the project source directory to execute the example notebook locally.
+
+This command will automatically handle the installation of necessary dependencies and execute the examples within a self-contained environment

{jaxsim-0.2.1.dev56 → jaxsim-0.2.1.dev69}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.1.dev56'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev56')
+__version__ = version = '0.2.1.dev69'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev69')

{jaxsim-0.2.1.dev56 → jaxsim-0.2.1.dev69}/src/jaxsim/parsers/descriptions/model.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import dataclasses
 import itertools
-from typing import List
+from typing import Sequence
 
 from jaxsim import logging
 
@@ -13,63 +15,62 @@ from .link import LinkDescription
 @dataclasses.dataclass(frozen=True)
 class ModelDescription(KinematicGraph):
     """
-    Description of a robotic model including links, joints, and collision shapes.
-
-    Args:
-        name (str): The name of the model.
-        fixed_base (bool): Indicates whether the model has a fixed base.
-        collision_shapes (List[CollisionShape]): List of collision shapes associated with the model.
+    Intermediate representation representing the kinematic graph of a robot model.
 
     Attributes:
-        name (str): The name of the model.
-        fixed_base (bool): Indicates whether the model has a fixed base.
-        collision_shapes (List[CollisionShape]): List of collision shapes associated with the model.
+        name: The name of the model.
+        fixed_base: Whether the model is either fixed-base or floating-base.
+        collision_shapes: List of collision shapes associated with the model.
     """
 
     name: str = None
+
     fixed_base: bool = True
-    collision_shapes: List[CollisionShape] = dataclasses.field(default_factory=list)
+
+    collision_shapes: list[CollisionShape] = dataclasses.field(
+        default_factory=list, repr=False, hash=False
+    )
 
     @staticmethod
     def build_model_from(
         name: str,
-        links: List[LinkDescription],
-        joints: List[JointDescription],
-        collisions: List[CollisionShape] = (),
+        links: list[LinkDescription],
+        joints: list[JointDescription],
+        frames: list[LinkDescription] | None = None,
+        collisions: list[CollisionShape] = (),
         fixed_base: bool = False,
         base_link_name: str | None = None,
-        considered_joints: List[str] | None = None,
+        considered_joints: Sequence[str] | None = None,
         model_pose: RootPose = RootPose(),
-    ) -> "ModelDescription":
+    ) -> ModelDescription:
         """
         Build a model description from provided components.
 
         Args:
-            name (str): The name of the model.
-            links (List[LinkDescription]): List of link descriptions.
-            joints (List[JointDescription]): List of joint descriptions.
-            collisions (List[CollisionShape]): List of collision shapes associated with the model.
-            fixed_base (bool): Indicates whether the model has a fixed base.
-            base_link_name (str): Name of the base link.
-            considered_joints (List[str]): List of joint names to consider.
-            model_pose (RootPose): Pose of the model's root.
+            name: The name of the model.
+            links: List of link descriptions.
+            joints: List of joint descriptions.
+            frames: List of frame descriptions.
+            collisions: List of collision shapes associated with the model.
+            fixed_base: Indicates whether the model has a fixed base.
+            base_link_name: Name of the base link (i.e. the root of the kinematic tree).
+            considered_joints: List of joint names to consider (by default all joints).
+            model_pose: Pose of the model's root (by default an identity transform).
 
         Returns:
-            ModelDescription: A ModelDescription instance representing the model.
-
-        Raises:
-            ValueError: If invalid or missing input data.
+            A ModelDescription instance representing the model.
         """
 
-        # Create the full kinematic graph
+        # Create the full kinematic graph.
         kinematic_graph = KinematicGraph.build_from(
             links=links,
             joints=joints,
+            frames=frames,
             root_link_name=base_link_name,
             root_pose=model_pose,
         )
 
-        # Reduce the graph if needed
+        # Reduce the graph if needed.
         if considered_joints is not None:
             kinematic_graph = kinematic_graph.reduce(
                 considered_joints=considered_joints
@@ -78,11 +79,13 @@ class ModelDescription(KinematicGraph):
         # Create the object to compute forward kinematics.
         fk = KinematicGraphTransforms(graph=kinematic_graph)
 
-        # Store here the final model collisions
-        final_collisions: List[CollisionShape] = []
+        # Container of the final model's collision shapes.
+        final_collisions: list[CollisionShape] = []
 
-        # Move and express the collision shapes of the removed link to the lumped link
+        # Move and express the collision shapes of removed links to the resulting
+        # lumped link that replace the combination of the removed link and its parent.
         for collision_shape in collisions:
+
             # Get all the collidable points of the shape
             coll_points = list(collision_shape.collidable_points)
 
@@ -112,7 +115,7 @@ class ModelDescription(KinematicGraph):
             final_collisions.append(new_collision_shape)
 
             # If the frame was found, update the collidable points' pose and add them
-            #  to the new collision shape
+            # to the new collision shape.
             for cp in collision_shape.collidable_points:
                 # Find the link that is part of the (reduced) model in which the
                 # collision shape's parent was lumped into
@@ -145,22 +148,20 @@ class ModelDescription(KinematicGraph):
             _joints_removed=kinematic_graph._joints_removed,
         )
 
+        # Check that the root link of kinematic graph is the desired base link.
         assert kinematic_graph.root.name == base_link_name, kinematic_graph.root.name
 
         return model
 
-    def reduce(self, considered_joints: List[str]) -> "ModelDescription":
+    def reduce(self, considered_joints: Sequence[str]) -> ModelDescription:
         """
         Reduce the model by removing specified joints.
 
         Args:
-            considered_joints (List[str]): List of joint names to consider.
+            The joint names to consider.
 
         Returns:
-            ModelDescription: A reduced ModelDescription instance.
-
-        Raises:
-            ValueError: If the specified joints are not part of the model.
+            A `ModelDescription` instance that only includes the considered joints.
         """
 
         if len(set(considered_joints) - set(self.joint_names())) != 0:
@@ -172,6 +173,7 @@ class ModelDescription(KinematicGraph):
             name=self.name,
             links=list(self.links_dict.values()),
             joints=self.joints,
+            frames=self.frames,
             collisions=self.collision_shapes,
             fixed_base=self.fixed_base,
             base_link_name=list(iter(self))[0].name,
@@ -190,12 +192,8 @@ class ModelDescription(KinematicGraph):
         Enable or disable collision shapes associated with a link.
 
         Args:
-            link_name (str): Name of the link.
-            enabled (bool): Enable or disable collision shapes associated with the link.
-
-        Raises:
-            ValueError: If the link name is not found in the model.
-
+            link_name: The name of the link.
+            enabled: Enable or disable collision shapes associated with the link.
         """
 
         if link_name not in self.link_names():
@@ -211,14 +209,10 @@ class ModelDescription(KinematicGraph):
         Get the collision shape associated with a specific link.
 
         Args:
-            link_name (str): Name of the link.
+            link_name: The name of the link.
 
         Returns:
-            CollisionShape: The collision shape associated with the link.
-
-        Raises:
-            ValueError: If the link name is not found in the model.
-
+            The collision shape associated with the link.
         """
 
         if link_name not in self.link_names():
@@ -233,14 +227,15 @@ class ModelDescription(KinematicGraph):
             ]
         )
 
-    def all_enabled_collidable_points(self) -> List[CollidablePoint]:
+    def all_enabled_collidable_points(self) -> list[CollidablePoint]:
         """
         Get all enabled collidable points in the model.
 
         Returns:
-            List[CollidablePoint]: A list of all enabled collidable points.
+            The list of all enabled collidable points.
 
         """
+
         # Get iterator of all collidable points
         all_collidable_points = itertools.chain.from_iterable(
             [shape.collidable_points for shape in self.collision_shapes]