From ec03aa9c9ba779811f4fd254d86f16a41e863e08 Mon Sep 17 00:00:00 2001 From: fanmin shi Date: Sun, 30 Jun 2024 23:11:26 +0000 Subject: [PATCH] Fix missing import in training_apg.ipynb When you run the notebook, you will immedidately observe ModuleNotFoundError for few modules. This installs required dependencies. --- mjx/training_apg.ipynb | 151 ++++++++++++++++++++++++++++++++--------- 1 file changed, 118 insertions(+), 33 deletions(-) diff --git a/mjx/training_apg.ipynb b/mjx/training_apg.ipynb index 26879392..4a7192dc 100644 --- a/mjx/training_apg.ipynb +++ b/mjx/training_apg.ipynb @@ -1,39 +1,39 @@ { "cells": [ { - "cell_type": "markdown", - "metadata": { - "id": "MpkYHwCqk7W-" - }, - "source": [ - "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n", - "\n", - "#

Tutorial

\n", - "\n", - "This notebook provides a tutorial for differentiable physics for policy learning in [**MuJoCo XLA (MJX)**](https://github.com/google-deepmind/mujoco/blob/main/mjx), a JAX-based implementation of MuJoCo.\n", - "\n", - "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n", - "\n", - "\n", - "This notebook was written by [Jing Yuan Luo](https://github.com/Andrew-Luo1).\n", - "\n", - "\u003c!-- Copyright 2021 DeepMind Technologies Limited\n", - "\n", - " Licensed under the Apache License, Version 2.0 (the \"License\");\n", - " you may not use this file except in compliance with the License.\n", - " You may obtain a copy of the License at\n", - "\n", - " http://www.apache.org/licenses/LICENSE-2.0\n", - "\n", - " Unless required by applicable law or agreed to in writing, software\n", - " distributed under the License is distributed on an \"AS IS\" BASIS,\n", - " WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", - " See the License for the specific language governing permissions and\n", - " limitations under the License.\n", - "--\u003e" - ] + "cell_type": "markdown", + "metadata": { + "id": "MpkYHwCqk7W-" }, - { + "source": [ + "![MuJoCo banner](https://raw.githubusercontent.com/google-deepmind/mujoco/main/banner.png)\n", + "\n", + "#

Tutorial

\n", + "\n", + "This notebook provides a tutorial for differentiable physics for policy learning in [**MuJoCo XLA (MJX)**](https://github.com/google-deepmind/mujoco/blob/main/mjx), a JAX-based implementation of MuJoCo.\n", + "\n", + "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n", + "\n", + "\n", + "This notebook was written by [Jing Yuan Luo](https://github.com/Andrew-Luo1).\n", + "\n", + "" + ] + }, + { "cell_type": "markdown", "metadata": {}, "source": [ @@ -106,12 +106,97 @@ "## Setup: Imports and installations" ] }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": {}, + "outputs": [], + "source": [ + "# Install MuJoCo, MJX, and Brax\n", + "!pip install mujoco\n", + "!pip install mujoco_mjx\n", + "!pip install brax" + ] + }, + { + "cell_type": "code", + "execution_count": 0, + "metadata": {}, + "outputs": [], + "source": [ + "#@title Check if MuJoCo installation was successful\n", + "\n", + "# Set up GPU rendering.\n", + "from google.colab import files\n", + "import distutils.util\n", + "import os\n", + "import subprocess\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", + "# 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", + " \"file_format_version\" : \"1.0.0\",\n", + " \"ICD\" : {\n", + " \"library_path\" : \"libEGL_nvidia.so.0\"\n", + " }\n", + "}\n", + "\"\"\")\n", + "\n", + "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n", + "print('Setting environment variable to use GPU rendering:')\n", + "%env MUJOCO_GL=egl\n", + "\n", + "# Check if installation was succesful.\n", + "try:\n", + " print('Checking that the installation succeeded:')\n", + " import mujoco\n", + " mujoco.MjModel.from_xml_string('')\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", + "print('Installation successful.')\n", + "\n", + "# Other imports and helper functions\n", + "import time\n", + "import itertools\n", + "import numpy as np\n", + "\n", + "# Graphics and plotting.\n", + "print('Installing mediapy:')\n", + "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n", + "!pip install -q mediapy\n", + "import mediapy as media\n", + "import matplotlib.pyplot as plt\n", + "\n", + "# More legible printing from numpy.\n", + "np.set_printoptions(precision=3, suppress=True, linewidth=100)\n", + "\n", + "from IPython.display import clear_output\n", + "clear_output()" + ] + }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ + "#@title Import MuJoCo, MJX, and Brax\n", + "\n", "import os\n", "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"0.8\" # 0.9 causes too much lag. \n", "from datetime import datetime\n", @@ -162,7 +247,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 0, "metadata": {}, "outputs": [], "source": [