Set DISPLAY environment variable for Colab

examples/PD_controller.ipynb

@@ -27,12 +27,22 @@
     "\n",
     "# Install JAX and Gazebo SDF\n",
     "if IS_COLAB:\n",
-    "    !{sys.executable} -m pip install -qU jaxsim\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",