diff --git a/examples/PD_controller.ipynb b/examples/PD_controller.ipynb index d15ef59cc..5da316c89 100644 --- a/examples/PD_controller.ipynb +++ b/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", + "\n", + " \"Open\n", + "" ] }, { @@ -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", diff --git a/examples/Parallel_computing.ipynb b/examples/Parallel_computing.ipynb index a3fc2ed30..4bb8fa9d9 100644 --- a/examples/Parallel_computing.ipynb +++ b/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", diff --git a/examples/README.md b/examples/README.md index 8fed061e7..e0303c8bb 100644 --- a/examples/README.md +++ b/examples/README.md @@ -4,15 +4,27 @@ This folder includes a Jupyter Notebook demonstrating the practical usage of JAX ### Examples -- [PD_controller](./PD_controller.ipynb) - 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) - An example demonstrating how to simulate vectorized models in parallel using JAXsim. +- [PD_controller](./PD_controller.ipynb) + Open In Colab + - 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) + Open In Colab + - 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, you can use [pixi](https://pixi.sh) to run the examples in a local environment: +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`:**