{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "atOI6yOnkxI6" }, "source": [ "# SimplerEnv: Simulated Manipulation Policy Evaluation for Real-World Robots\n", "\n", "- Project page: \n", "- Code: " ] }, { "cell_type": "markdown", "metadata": { "id": "3E3ILp1cTULl" }, "source": [ "## Installation\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "T5aYcrQzH9qo", "outputId": "7ff77e31-2cc3-4500-de1e-4e396636f33b" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Fri Aug 9 22:59:30 2024 \n", "+---------------------------------------------------------------------------------------+\n", "| NVIDIA-SMI 535.104.05 Driver Version: 535.104.05 CUDA Version: 12.2 |\n", "|-----------------------------------------+----------------------+----------------------+\n", "| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |\n", "| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |\n", "| | | MIG M. |\n", "|=========================================+======================+======================|\n", "| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |\n", "| N/A 46C P8 9W / 70W | 0MiB / 15360MiB | 0% Default |\n", "| | | N/A |\n", "+-----------------------------------------+----------------------+----------------------+\n", " \n", "+---------------------------------------------------------------------------------------+\n", "| Processes: |\n", "| GPU GI CI PID Type Process name GPU Memory |\n", "| ID ID Usage |\n", "|=======================================================================================|\n", "| No running processes found |\n", "+---------------------------------------------------------------------------------------+\n" ] } ], "source": [ "#@title [!Important]Please use a GPU runtime.\n", "!nvidia-smi" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "id": "pI0Ge9viDPks" }, "outputs": [], "source": [ "# @title Install vulkan for rendering\n", "!apt-get install -yqq --no-install-recommends libvulkan-dev vulkan-tools\n", "# below fixes some bugs introduced by some recent Colab changes\n", "!mkdir -p /usr/share/vulkan/icd.d\n", "!wget -q -P /usr/share/vulkan/icd.d https://raw.githubusercontent.com/haosulab/ManiSkill/main/docker/nvidia_icd.json\n", "!wget -q -O /usr/share/glvnd/egl_vendor.d/10_nvidia.json https://raw.githubusercontent.com/haosulab/ManiSkill/main/docker/10_nvidia.json" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "r_EBPnGtU1zA", "outputId": "8f5558c9-3365-4d8f-f259-8773ee56a434" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "'DISPLAY' environment variable not set... skipping surface info\n", "error: XDG_RUNTIME_DIR not set in the environment.\n", "==========\n", "VULKANINFO\n", "==========\n", "\n", "Vulkan Instance Version: 1.3.204\n" ] } ], "source": [ "# @title Make sure vulkan is installed correctly\n", "!vulkaninfo | head -n 5" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "ZFUisRykZ4l4", "outputId": "fbdb2769-4798-45f0-e835-39d98a02a392" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Requirement already satisfied: numpy==1.24.4 in /usr/local/lib/python3.10/dist-packages (1.24.4)\n", "Requirement already satisfied: orbax-checkpoint==0.4.4 in /usr/local/lib/python3.10/dist-packages (0.4.4)\n", "Requirement already satisfied: absl-py in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (1.4.0)\n", "Requirement already satisfied: etils[epath,epy] in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (1.7.0)\n", "Requirement already satisfied: typing_extensions in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (4.5.0)\n", "Requirement already satisfied: msgpack in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (1.0.8)\n", "Requirement already satisfied: jax>=0.4.9 in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (0.4.20)\n", "Requirement already satisfied: jaxlib in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (0.4.20+cuda12.cudnn89)\n", "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (1.24.4)\n", "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (6.0.2)\n", "Requirement already satisfied: tensorstore>=0.1.35 in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (0.1.64)\n", "Requirement already satisfied: nest_asyncio in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (1.6.0)\n", "Requirement already satisfied: protobuf in /usr/local/lib/python3.10/dist-packages (from orbax-checkpoint==0.4.4) (3.20.3)\n", "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.9->orbax-checkpoint==0.4.4) (0.3.2)\n", "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.9->orbax-checkpoint==0.4.4) (3.3.0)\n", "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax>=0.4.9->orbax-checkpoint==0.4.4) (1.13.1)\n", "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from etils[epath,epy]->orbax-checkpoint==0.4.4) (2024.6.1)\n", "Requirement already satisfied: importlib_resources in /usr/local/lib/python3.10/dist-packages (from etils[epath,epy]->orbax-checkpoint==0.4.4) (6.4.0)\n", "Requirement already satisfied: zipp in /usr/local/lib/python3.10/dist-packages (from etils[epath,epy]->orbax-checkpoint==0.4.4) (3.19.2)\n", "Collecting scipy==1.12.0\n", " Downloading scipy-1.12.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (60 kB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.4/60.4 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hRequirement already satisfied: numpy<1.29.0,>=1.22.4 in /usr/local/lib/python3.10/dist-packages (from scipy==1.12.0) (1.24.4)\n", "Downloading scipy-1.12.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (38.4 MB)\n", "\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m38.4/38.4 MB\u001b[0m \u001b[31m36.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n", "\u001b[?25hInstalling collected packages: scipy\n", " Attempting uninstall: scipy\n", " Found existing installation: scipy 1.13.1\n", " Uninstalling scipy-1.13.1:\n", " Successfully uninstalled scipy-1.13.1\n", "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n", "xgboost 2.1.1 requires nvidia-nccl-cu12; platform_system == \"Linux\" and platform_machine != \"aarch64\", which is not installed.\n", "albumentations 1.4.13 requires typing-extensions>=4.9.0, but you have typing-extensions 4.5.0 which is incompatible.\n", "osqp 0.6.7.post0 requires scipy!=1.12.0,>=0.13.2, but you have scipy 1.12.0 which is incompatible.\u001b[0m\u001b[31m\n", "\u001b[0mSuccessfully installed scipy-1.12.0\n", "Requirement already satisfied: keras==2.15.0 in /usr/local/lib/python3.10/dist-packages (2.15.0)\n", "Requirement already satisfied: tensorflow==2.15.1 in /usr/local/lib/python3.10/dist-packages (2.15.1)\n", "Requirement already satisfied: absl-py>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (1.4.0)\n", "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (1.6.3)\n", "Requirement already satisfied: flatbuffers>=23.5.26 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (24.3.25)\n", "Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (0.6.0)\n", "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (0.2.0)\n", "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (3.11.0)\n", "Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (18.1.1)\n", "Requirement already satisfied: ml-dtypes~=0.3.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (0.3.2)\n", "Requirement already satisfied: numpy<2.0.0,>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (1.24.4)\n", "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (3.3.0)\n", "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (24.1)\n", "Requirement already satisfied: protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (3.20.3)\n", "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (71.0.4)\n", "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (1.16.0)\n", "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (2.4.0)\n", "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (4.5.0)\n", "Requirement already satisfied: wrapt<1.15,>=1.11.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (1.14.1)\n", "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (0.37.1)\n", "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (1.64.1)\n", "Requirement already satisfied: tensorboard<2.16,>=2.15 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (2.15.2)\n", "Requirement already satisfied: tensorflow-estimator<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (2.15.0)\n", "Requirement already satisfied: keras<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.1) (2.15.0)\n", "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from astunparse>=1.6.0->tensorflow==2.15.1) (0.44.0)\n", "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.1) (2.27.0)\n", "Requirement already satisfied: google-auth-oauthlib<2,>=0.5 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.1) (1.2.1)\n", "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.1) (3.6)\n", "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.1) (2.32.3)\n", "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.1) (0.7.2)\n", "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.1) (3.0.3)\n", "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (5.4.0)\n", "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (0.4.0)\n", "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (4.9)\n", "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (1.3.1)\n", "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (3.3.2)\n", "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (3.7)\n", "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (2.0.7)\n", "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3,>=2.21.0->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (2024.7.4)\n", "Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from werkzeug>=1.0.1->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (2.1.5)\n", "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (0.6.0)\n", "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.10/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow==2.15.1) (3.2.2)\n", "fatal: destination path 'ManiSkill2_real2sim' already exists and is not an empty directory.\n", "Obtaining file:///content/ManiSkill2_real2sim\n", " Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n", " Checking if build backend supports build_editable ... \u001b[?25l\u001b[?25hdone\n", " Getting requirements to build editable ... \u001b[?25l\u001b[?25hdone\n", " Preparing editable metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (1.24.4)\n", "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (1.12.0)\n", "Requirement already satisfied: gymnasium>=0.28.1 in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (0.29.1)\n", "Requirement already satisfied: sapien==2.2.2 in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (2.2.2)\n", "Requirement already satisfied: h5py in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (3.11.0)\n", "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (6.0.2)\n", "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (4.66.5)\n", "Requirement already satisfied: GitPython in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (3.1.43)\n", "Requirement already satisfied: tabulate in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (0.9.0)\n", "Requirement already satisfied: gdown>=4.6.0 in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (5.1.0)\n", "Requirement already satisfied: transforms3d in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (0.4.2)\n", "Requirement already satisfied: opencv-python in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (4.10.0.84)\n", "Requirement already satisfied: imageio in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (2.34.2)\n", "Requirement already satisfied: trimesh in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (4.4.4)\n", "Requirement already satisfied: rtree in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (1.3.0)\n", "Requirement already satisfied: ruckig in /usr/local/lib/python3.10/dist-packages (from mani_skill2_real2sim==0.5.3) (0.12.2)\n", "Requirement already satisfied: requests>=2.22 in /usr/local/lib/python3.10/dist-packages (from sapien==2.2.2->mani_skill2_real2sim==0.5.3) (2.32.3)\n", "Requirement already satisfied: beautifulsoup4 in /usr/local/lib/python3.10/dist-packages (from gdown>=4.6.0->mani_skill2_real2sim==0.5.3) (4.12.3)\n", "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from gdown>=4.6.0->mani_skill2_real2sim==0.5.3) (3.15.4)\n", "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.10/dist-packages (from gymnasium>=0.28.1->mani_skill2_real2sim==0.5.3) (2.2.1)\n", "Requirement already satisfied: typing-extensions>=4.3.0 in /usr/local/lib/python3.10/dist-packages (from gymnasium>=0.28.1->mani_skill2_real2sim==0.5.3) (4.5.0)\n", "Requirement already satisfied: farama-notifications>=0.0.1 in /usr/local/lib/python3.10/dist-packages (from gymnasium>=0.28.1->mani_skill2_real2sim==0.5.3) (0.0.4)\n", "Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.10/dist-packages (from GitPython->mani_skill2_real2sim==0.5.3) (4.0.11)\n", "Requirement already satisfied: pillow>=8.3.2 in /usr/local/lib/python3.10/dist-packages (from imageio->mani_skill2_real2sim==0.5.3) (9.4.0)\n", "Requirement already satisfied: imageio-ffmpeg in /usr/local/lib/python3.10/dist-packages (from imageio[ffmpeg]->mani_skill2_real2sim==0.5.3) (0.5.1)\n", "Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from imageio[ffmpeg]->mani_skill2_real2sim==0.5.3) (5.9.5)\n", "Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.10/dist-packages (from gitdb<5,>=4.0.1->GitPython->mani_skill2_real2sim==0.5.3) (5.0.1)\n", "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.22->sapien==2.2.2->mani_skill2_real2sim==0.5.3) (3.3.2)\n", "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.22->sapien==2.2.2->mani_skill2_real2sim==0.5.3) (3.7)\n", "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.22->sapien==2.2.2->mani_skill2_real2sim==0.5.3) (2.0.7)\n", "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.22->sapien==2.2.2->mani_skill2_real2sim==0.5.3) (2024.7.4)\n", "Requirement already satisfied: soupsieve>1.2 in /usr/local/lib/python3.10/dist-packages (from beautifulsoup4->gdown>=4.6.0->mani_skill2_real2sim==0.5.3) (2.5)\n", "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from imageio-ffmpeg->imageio[ffmpeg]->mani_skill2_real2sim==0.5.3) (71.0.4)\n", "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /usr/local/lib/python3.10/dist-packages (from requests[socks]->gdown>=4.6.0->mani_skill2_real2sim==0.5.3) (1.7.1)\n", "Building wheels for collected packages: mani_skill2_real2sim\n", " Building editable for mani_skill2_real2sim (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", " Created wheel for mani_skill2_real2sim: filename=mani_skill2_real2sim-0.5.3-0.editable-py3-none-any.whl size=7553 sha256=7301927181a4813c8736b8599614213964c84b4d63c3d70bb5914a7c3f180f76\n", " Stored in directory: /tmp/pip-ephem-wheel-cache-q405ysvq/wheels/e8/f9/5c/bae03205d26a1eae11ad649a4eea25c9ec62eec90dfa4d738c\n", "Successfully built mani_skill2_real2sim\n", "Installing collected packages: mani_skill2_real2sim\n", " Attempting uninstall: mani_skill2_real2sim\n", " Found existing installation: mani_skill2_real2sim 0.5.3\n", " Uninstalling mani_skill2_real2sim-0.5.3:\n", " Successfully uninstalled mani_skill2_real2sim-0.5.3\n", "Successfully installed mani_skill2_real2sim-0.5.3\n", "fatal: destination path 'SimplerEnv' already exists and is not an empty directory.\n", "Obtaining file:///content/SimplerEnv\n", " Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n", " Checking if build backend supports build_editable ... \u001b[?25l\u001b[?25hdone\n", " Getting requirements to build editable ... \u001b[?25l\u001b[?25hdone\n", " Preparing editable metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", "Building wheels for collected packages: simpler_env\n", " Building editable for simpler_env (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", " Created wheel for simpler_env: filename=simpler_env-0.0.1-0.editable-py3-none-any.whl size=3558 sha256=11911f9ddf571b4aea70e40e18e5703be9042dcaccc1decb47d8285e00d61d05\n", " Stored in directory: /tmp/pip-ephem-wheel-cache-580g7qm8/wheels/35/70/91/7fd0a26ad5a9a0a798cd6c8ee196b326da951fa1853ca5a8b3\n", "Successfully built simpler_env\n", "Installing collected packages: simpler_env\n", " Attempting uninstall: simpler_env\n", " Found existing installation: simpler_env 0.0.1\n", " Uninstalling simpler_env-0.0.1:\n", " Successfully uninstalled simpler_env-0.0.1\n", "Successfully installed simpler_env-0.0.1\n", "mkdir: cannot create directory ‘./SimplerEnv/checkpoints’: File exists\n" ] } ], "source": [ "# @title Install Real2Sim\n", "!pip install numpy==1.24.4\n", "!pip install orbax-checkpoint==0.4.4\n", "!pip install scipy==1.12.0\n", "!pip install keras==2.15.0\n", "!pip install tensorflow==2.15.1\n", "!git clone https://github.com/simpler-env/ManiSkill2_real2sim.git\n", "!pip install -e ./ManiSkill2_real2sim\n", "!git clone https://github.com/simpler-env/SimplerEnv.git\n", "!pip install -e ./SimplerEnv\n", "!mkdir ./SimplerEnv/checkpoints" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "id": "MMuGj8uqmbgb" }, "outputs": [], "source": [ "#@title [Optional]Install RT-1 dependencies\n", "! pip install --quiet tf_agents" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "Ja1xGJe5mnuu", "outputId": "d38cbb2d-b9fe-4cd7-97ad-56f6b1905cfb" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "fatal: destination path 'octo' already exists and is not an empty directory.\n", "HEAD is now at 653c54a Merge pull request #31 from octo-models/fix_robot_eval\n", "Obtaining file:///content/octo\n", " Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n", " Checking if build backend supports build_editable ... \u001b[?25l\u001b[?25hdone\n", " Getting requirements to build editable ... \u001b[?25l\u001b[?25hdone\n", " Preparing editable metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", "Building wheels for collected packages: octo\n", " Building editable for octo (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n", " Created wheel for octo: filename=octo-0.0.0-0.editable-py3-none-any.whl size=3382 sha256=0883671cae31e9dff45138a036864480dd53c96221e3659aebc0e49f66302052\n", " Stored in directory: /tmp/pip-ephem-wheel-cache-roxdztlh/wheels/66/42/18/cf44976e18724f184efd64bbc31830460087187fa524fa6e42\n", "Successfully built octo\n", "Installing collected packages: octo\n", " Attempting uninstall: octo\n", " Found existing installation: octo 0.0.0\n", " Uninstalling octo-0.0.0:\n", " Successfully uninstalled octo-0.0.0\n", "Successfully installed octo-0.0.0\n", "Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n", "Requirement already satisfied: jax==0.4.20 in /usr/local/lib/python3.10/dist-packages (from jax[cuda12_pip]==0.4.20) (0.4.20)\n", "Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n", "Requirement already satisfied: jaxlib==0.4.20 in /usr/local/lib/python3.10/dist-packages (from jaxlib[cuda12_pip]==0.4.20) (0.4.20+cuda12.cudnn89)\n", "Requirement already satisfied: jaxlib==0.4.20+cuda12.cudnn89 in /usr/local/lib/python3.10/dist-packages (0.4.20+cuda12.cudnn89)\n", "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jaxlib==0.4.20+cuda12.cudnn89) (1.12.0)\n", "Requirement already satisfied: numpy>=1.22 in /usr/local/lib/python3.10/dist-packages (from jaxlib==0.4.20+cuda12.cudnn89) (1.24.4)\n", "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jaxlib==0.4.20+cuda12.cudnn89) (0.3.2)\n", "Requirement already satisfied: distrax==0.1.5 in /usr/local/lib/python3.10/dist-packages (0.1.5)\n", "Requirement already satisfied: einops>=0.6.1 in /usr/local/lib/python3.10/dist-packages (0.8.0)\n", "Requirement already satisfied: absl-py>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from distrax==0.1.5) (1.4.0)\n", "Requirement already satisfied: chex>=0.1.8 in /usr/local/lib/python3.10/dist-packages (from distrax==0.1.5) (0.1.86)\n", "Requirement already satisfied: jax>=0.1.55 in /usr/local/lib/python3.10/dist-packages (from distrax==0.1.5) (0.4.20)\n", "Requirement already satisfied: jaxlib>=0.1.67 in /usr/local/lib/python3.10/dist-packages (from distrax==0.1.5) (0.4.20+cuda12.cudnn89)\n", "Requirement already satisfied: numpy>=1.23.0 in /usr/local/lib/python3.10/dist-packages (from distrax==0.1.5) (1.24.4)\n", "Requirement already satisfied: tensorflow-probability>=0.15.0 in /usr/local/lib/python3.10/dist-packages (from distrax==0.1.5) (0.23.0)\n", "Requirement already satisfied: typing-extensions>=4.2.0 in /usr/local/lib/python3.10/dist-packages (from chex>=0.1.8->distrax==0.1.5) (4.5.0)\n", "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from chex>=0.1.8->distrax==0.1.5) (0.12.1)\n", "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax>=0.1.55->distrax==0.1.5) (0.3.2)\n", "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax>=0.1.55->distrax==0.1.5) (3.3.0)\n", "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax>=0.1.55->distrax==0.1.5) (1.12.0)\n", "Requirement already satisfied: six>=1.10.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability>=0.15.0->distrax==0.1.5) (1.16.0)\n", "Requirement already satisfied: decorator in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability>=0.15.0->distrax==0.1.5) (4.4.2)\n", "Requirement already satisfied: cloudpickle>=1.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability>=0.15.0->distrax==0.1.5) (2.2.1)\n", "Requirement already satisfied: gast>=0.3.2 in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability>=0.15.0->distrax==0.1.5) (0.6.0)\n", "Requirement already satisfied: dm-tree in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability>=0.15.0->distrax==0.1.5) (0.1.8)\n" ] } ], "source": [ "#@title [Optional]Install Octo dependencies\n", "!git clone https://github.com/octo-models/octo\n", "!cd ./octo && git checkout 653c54acde686fde619855f2eac0dd6edad7116b && cd ..\n", "!pip install -e ./octo\n", "!pip install --upgrade \"jax[cuda12_pip]==0.4.20\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --no-deps\n", "!pip install --upgrade \"jaxlib[cuda12_pip]==0.4.20\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --no-deps\n", "!pip install jaxlib==0.4.20+cuda12.cudnn89\n", "!pip install distrax==0.1.5 \"einops>= 0.6.1\"" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "id": "tSGomwvFcKf-" }, "outputs": [], "source": [ "# @title Install other requirements\n", "!pip install --quiet mediapy" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "id": "4stLX8cdaWD3" }, "outputs": [], "source": [ "# @title [Important]Post Installation\n", "\n", "# run this so local pip installs are recognized\n", "import site\n", "site.main()" ] }, { "cell_type": "markdown", "metadata": { "id": "TZJMaVVObgRt" }, "source": [ "## Create a Simulated Environment and Take Random Actions" ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 682 }, "id": "pMhxVs4nC9NT", "outputId": "fd0005a4-d0aa-42e8-b739-de266edb8916" }, "outputs": [ { "output_type": "stream", "name": "stderr", "text": [ "/usr/local/lib/python3.10/dist-packages/gymnasium/core.py:311: UserWarning: \u001b[33mWARN: env.get_language_instruction to get variables from other wrappers is deprecated and will be removed in v1.0, to get this variable you can do `env.unwrapped.get_language_instruction` for environment variables or `env.get_wrapper_attr('get_language_instruction')` that will search the reminding wrappers.\u001b[0m\n", " logger.warn(\n", "/usr/local/lib/python3.10/dist-packages/gymnasium/core.py:311: UserWarning: \u001b[33mWARN: env.robot_uid to get variables from other wrappers is deprecated and will be removed in v1.0, to get this variable you can do `env.unwrapped.robot_uid` for environment variables or `env.get_wrapper_attr('robot_uid')` that will search the reminding wrappers.\u001b[0m\n", " logger.warn(\n" ] }, { "output_type": "stream", "name": "stdout", "text": [ "Reset info {'scene_name': 'google_pick_coke_can_1_v4', 'scene_offset': None, 'scene_pose': None, 'scene_table_height': 0.87, 'urdf_version': 'recolor_tabletop_visual_matching_1', 'rgb_overlay_path': '/content/ManiSkill2_real2sim/data/real_inpainting/google_coke_can_real_eval_1.png', 'rgb_overlay_cameras': ['overhead_camera'], 'rgb_overlay_mode': 'background', 'disable_bad_material': False, 'model_id': 'opened_coke_can', 'model_scale': 1.0, 'distractor_model_ids': None, 'distractor_model_scales': None, 'obj_init_pose_wrt_robot_base': Pose([0.587925, -0.0238302, 0.840576], [0.707052, -0.0081018, -0.01162, -0.70702]), 'orientation': 'laid_vertically'}\n", "Instruction pick coke can\n", "Episode stats OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])\n" ] }, { "output_type": "display_data", "data": { "text/plain": [ "" ], "text/html": [ "
" ] }, "metadata": {} } ], "source": [ "import simpler_env\n", "from simpler_env.utils.env.observation_utils import get_image_from_maniskill2_obs_dict\n", "import mediapy\n", "import sapien.core as sapien\n", "\n", "task_name = \"google_robot_pick_coke_can\" # @param [\"google_robot_pick_coke_can\", \"google_robot_move_near\", \"google_robot_open_drawer\", \"google_robot_close_drawer\", \"widowx_spoon_on_towel\", \"widowx_carrot_on_plate\", \"widowx_stack_cube\", \"widowx_put_eggplant_in_basket\"]\n", "\n", "if 'env' in locals():\n", " print(\"Closing existing env\")\n", " env.close()\n", " del env\n", "env = simpler_env.make(task_name)\n", "# Colab GPU does not supoort denoiser\n", "sapien.render_config.rt_use_denoiser = False\n", "obs, reset_info = env.reset()\n", "instruction = env.get_language_instruction()\n", "print(\"Reset info\", reset_info)\n", "print(\"Instruction\", instruction)\n", "\n", "frames = []\n", "done, truncated = False, False\n", "while not (done or truncated):\n", " # action[:3]: delta xyz; action[3:6]: delta rotation in axis-angle representation;\n", " # action[6:7]: gripper (the meaning of open / close depends on robot URDF)\n", " image = get_image_from_maniskill2_obs_dict(env, obs)\n", " action = env.action_space.sample() # replace this with your policy inference\n", " obs, reward, done, truncated, info = env.step(action)\n", " frames.append(image)\n", "\n", "episode_stats = info.get('episode_stats', {})\n", "print(\"Episode stats\", episode_stats)\n", "mediapy.show_video(frames, fps=10)" ] }, { "cell_type": "markdown", "metadata": { "id": "0Epp91kpcqWv" }, "source": [ "## Run Inference on Simulated Environments" ] }, { "cell_type": "code", "execution_count": 10, "metadata": { "id": "XgxjPOtxoAw6" }, "outputs": [], "source": [ "# @title Setup\n", "\n", "import os\n", "import numpy as np\n", "import simpler_env\n", "from simpler_env.utils.env.observation_utils import get_image_from_maniskill2_obs_dict\n", "import mediapy\n", "\n", "\n", "RT_1_CHECKPOINTS = {\n", " \"rt_1_x\": \"rt_1_x_tf_trained_for_002272480_step\",\n", " \"rt_1_400k\": \"rt_1_tf_trained_for_000400120\",\n", " \"rt_1_58k\": \"rt_1_tf_trained_for_000058240\",\n", " \"rt_1_1k\": \"rt_1_tf_trained_for_000001120\",\n", "}\n", "\n", "\n", "def get_rt_1_checkpoint(name, ckpt_dir=\"./SimplerEnv/checkpoints\"):\n", " assert name in RT_1_CHECKPOINTS, name\n", " ckpt_name = RT_1_CHECKPOINTS[name]\n", " ckpt_path = os.path.join(ckpt_dir, ckpt_name)\n", " if not os.path.exists(ckpt_path):\n", " if name == \"rt_1_x\":\n", " !gsutil -m cp -r gs://gdm-robotics-open-x-embodiment/open_x_embodiment_and_rt_x_oss/{ckpt_name}.zip {ckpt_dir}\n", " !unzip {ckpt_dir}/{ckpt_name}.zip -d {ckpt_dir}\n", " else:\n", " !gsutil -m cp -r gs://gdm-robotics-open-x-embodiment/open_x_embodiment_and_rt_x_oss/{ckpt_name} {ckpt_dir}\n", " return ckpt_path" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "p4IwfAe_rgDr", "outputId": "a3ef17ad-bb1d-412a-c741-c06791861725" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Closing existing env\n", "Reset info {'scene_name': 'google_pick_coke_can_1_v4', 'scene_offset': None, 'scene_pose': None, 'scene_table_height': 0.87, 'urdf_version': 'recolor_tabletop_visual_matching_1', 'rgb_overlay_path': '/content/ManiSkill2_real2sim/data/real_inpainting/google_coke_can_real_eval_1.png', 'rgb_overlay_cameras': ['overhead_camera'], 'rgb_overlay_mode': 'background', 'disable_bad_material': False, 'model_id': 'opened_coke_can', 'model_scale': 1.0, 'distractor_model_ids': None, 'distractor_model_scales': None, 'obj_init_pose_wrt_robot_base': Pose([0.587925, -0.0238302, 0.840576], [0.707052, -0.0081018, -0.01162, -0.70702]), 'orientation': 'laid_vertically'}\n", "Instruction pick coke can\n" ] } ], "source": [ "# @title Select your model and environment\n", "\n", "task_name = \"google_robot_pick_coke_can\" # @param [\"google_robot_pick_coke_can\", \"google_robot_move_near\", \"google_robot_open_drawer\", \"google_robot_close_drawer\", \"widowx_spoon_on_towel\", \"widowx_carrot_on_plate\", \"widowx_stack_cube\", \"widowx_put_eggplant_in_basket\"]\n", "\n", "if 'env' in locals():\n", " print(\"Closing existing env\")\n", " env.close()\n", " del env\n", "env = simpler_env.make(task_name)\n", "\n", "# Note: we turned off the denoiser as the colab kernel will crash if it's turned on\n", "# To use the denoiser, please git clone our SIMPLER environments\n", "# and perform evaluations locally.\n", "sapien.render_config.rt_use_denoiser = False\n", "\n", "obs, reset_info = env.reset()\n", "instruction = env.get_language_instruction()\n", "print(\"Reset info\", reset_info)\n", "print(\"Instruction\", instruction)\n", "\n", "if \"google\" in task_name:\n", " policy_setup = \"google_robot\"\n", "else:\n", " policy_setup = \"widowx_bridge\"" ] }, { "cell_type": "code", "execution_count": 14, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "J8kyi5Czp0yH", "outputId": "61f1c430-68d5-4a0d-efc6-c066f8f6c5af" }, "outputs": [ { "output_type": "stream", "name": "stderr", "text": [ "WARNING:absl:Importing a function (__inference_block4e_expand_activation_layer_call_and_return_conditional_losses_1756302) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7b_activation_layer_call_and_return_conditional_losses_1743046) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6c_se_reduce_layer_call_and_return_conditional_losses_1760957) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6f_expand_activation_layer_call_and_return_conditional_losses_1742603) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4e_se_reduce_layer_call_and_return_conditional_losses_1740331) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6b_activation_layer_call_and_return_conditional_losses_1760331) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4d_expand_activation_layer_call_and_return_conditional_losses_1755734) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4a_se_reduce_layer_call_and_return_conditional_losses_1739487) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6e_expand_activation_layer_call_and_return_conditional_losses_1742385) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6c_expand_activation_layer_call_and_return_conditional_losses_1741949) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5c_se_reduce_layer_call_and_return_conditional_losses_1758153) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5a_se_reduce_layer_call_and_return_conditional_losses_1757064) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2c_expand_activation_layer_call_and_return_conditional_losses_1751830) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3b_activation_layer_call_and_return_conditional_losses_1753066) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6a_activation_layer_call_and_return_conditional_losses_1741576) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5b_activation_layer_call_and_return_conditional_losses_1757527) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4a_se_reduce_layer_call_and_return_conditional_losses_1754271) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5d_se_reduce_layer_call_and_return_conditional_losses_1758721) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7a_activation_layer_call_and_return_conditional_losses_1763171) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5e_expand_activation_layer_call_and_return_conditional_losses_1741322) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3a_se_reduce_layer_call_and_return_conditional_losses_1752603) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4b_se_reduce_layer_call_and_return_conditional_losses_1739677) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5d_se_reduce_layer_call_and_return_conditional_losses_1741175) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1a_se_reduce_layer_call_and_return_conditional_losses_1737857) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4c_se_reduce_layer_call_and_return_conditional_losses_1755360) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5a_activation_layer_call_and_return_conditional_losses_1757006) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4c_activation_layer_call_and_return_conditional_losses_1739859) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1b_se_reduce_layer_call_and_return_conditional_losses_1738014) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4d_expand_activation_layer_call_and_return_conditional_losses_1740042) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1b_se_reduce_layer_call_and_return_conditional_losses_1750356) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3b_expand_activation_layer_call_and_return_conditional_losses_1752930) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6c_activation_layer_call_and_return_conditional_losses_1741984) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6c_se_reduce_layer_call_and_return_conditional_losses_1742020) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5d_expand_activation_layer_call_and_return_conditional_losses_1758527) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5a_se_reduce_layer_call_and_return_conditional_losses_1740549) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5b_expand_activation_layer_call_and_return_conditional_losses_1740668) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3b_expand_activation_layer_call_and_return_conditional_losses_1738979) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6a_activation_layer_call_and_return_conditional_losses_1759810) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6c_activation_layer_call_and_return_conditional_losses_1760899) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3c_se_reduce_layer_call_and_return_conditional_losses_1753692) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6a_se_reduce_layer_call_and_return_conditional_losses_1759868) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6d_expand_activation_layer_call_and_return_conditional_losses_1761331) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2a_se_reduce_layer_call_and_return_conditional_losses_1750935) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_top_activation_layer_call_and_return_conditional_losses_1743229) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2a_activation_layer_call_and_return_conditional_losses_1738197) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4b_expand_activation_layer_call_and_return_conditional_losses_1754598) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3c_expand_activation_layer_call_and_return_conditional_losses_1753498) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5e_se_reduce_layer_call_and_return_conditional_losses_1759289) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7b_activation_layer_call_and_return_conditional_losses_1763692) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6f_expand_activation_layer_call_and_return_conditional_losses_1762467) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2a_se_reduce_layer_call_and_return_conditional_losses_1738233) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2c_se_reduce_layer_call_and_return_conditional_losses_1752024) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5c_activation_layer_call_and_return_conditional_losses_1758095) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4d_activation_layer_call_and_return_conditional_losses_1740077) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5d_activation_layer_call_and_return_conditional_losses_1758663) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6f_activation_layer_call_and_return_conditional_losses_1762603) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5b_se_reduce_layer_call_and_return_conditional_losses_1740739) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6e_activation_layer_call_and_return_conditional_losses_1762035) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6c_expand_activation_layer_call_and_return_conditional_losses_1760763) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6b_activation_layer_call_and_return_conditional_losses_1741766) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6e_activation_layer_call_and_return_conditional_losses_1742420) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6f_activation_layer_call_and_return_conditional_losses_1742638) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5e_expand_activation_layer_call_and_return_conditional_losses_1759095) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6f_se_reduce_layer_call_and_return_conditional_losses_1742674) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5b_se_reduce_layer_call_and_return_conditional_losses_1757585) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4e_activation_layer_call_and_return_conditional_losses_1740295) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3b_activation_layer_call_and_return_conditional_losses_1739014) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5a_activation_layer_call_and_return_conditional_losses_1740513) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3b_se_reduce_layer_call_and_return_conditional_losses_1753124) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1a_se_reduce_layer_call_and_return_conditional_losses_1749969) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5c_expand_activation_layer_call_and_return_conditional_losses_1740886) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4e_expand_activation_layer_call_and_return_conditional_losses_1740260) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3c_se_reduce_layer_call_and_return_conditional_losses_1739268) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3a_activation_layer_call_and_return_conditional_losses_1752545) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6b_se_reduce_layer_call_and_return_conditional_losses_1741802) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1b_activation_layer_call_and_return_conditional_losses_1750298) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3a_expand_activation_layer_call_and_return_conditional_losses_1738788) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6d_activation_layer_call_and_return_conditional_losses_1742202) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3c_expand_activation_layer_call_and_return_conditional_losses_1739197) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5e_activation_layer_call_and_return_conditional_losses_1741357) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4b_expand_activation_layer_call_and_return_conditional_losses_1739606) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2c_se_reduce_layer_call_and_return_conditional_losses_1738641) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7b_expand_activation_layer_call_and_return_conditional_losses_1763556) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1a_activation_layer_call_and_return_conditional_losses_1749911) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7a_expand_activation_layer_call_and_return_conditional_losses_1763035) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5e_activation_layer_call_and_return_conditional_losses_1759231) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6a_expand_activation_layer_call_and_return_conditional_losses_1741540) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6d_se_reduce_layer_call_and_return_conditional_losses_1761525) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2a_expand_activation_layer_call_and_return_conditional_losses_1750730) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3a_expand_activation_layer_call_and_return_conditional_losses_1752398) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4a_activation_layer_call_and_return_conditional_losses_1739451) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_stem_activation_layer_call_and_return_conditional_losses_1749775) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5c_se_reduce_layer_call_and_return_conditional_losses_1740957) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4c_expand_activation_layer_call_and_return_conditional_losses_1739824) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4a_expand_activation_layer_call_and_return_conditional_losses_1739415) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5e_se_reduce_layer_call_and_return_conditional_losses_1741393) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3b_se_reduce_layer_call_and_return_conditional_losses_1739050) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5c_expand_activation_layer_call_and_return_conditional_losses_1757959) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4d_activation_layer_call_and_return_conditional_losses_1755870) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5b_expand_activation_layer_call_and_return_conditional_losses_1757391) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6b_expand_activation_layer_call_and_return_conditional_losses_1741731) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_stem_activation_layer_call_and_return_conditional_losses_1737786) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4e_se_reduce_layer_call_and_return_conditional_losses_1756496) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7a_activation_layer_call_and_return_conditional_losses_1742856) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5d_activation_layer_call_and_return_conditional_losses_1741139) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7a_se_reduce_layer_call_and_return_conditional_losses_1742892) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7b_se_reduce_layer_call_and_return_conditional_losses_1743082) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4e_activation_layer_call_and_return_conditional_losses_1756438) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4b_activation_layer_call_and_return_conditional_losses_1754734) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2b_expand_activation_layer_call_and_return_conditional_losses_1738352) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5d_expand_activation_layer_call_and_return_conditional_losses_1741104) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6b_se_reduce_layer_call_and_return_conditional_losses_1760389) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3c_activation_layer_call_and_return_conditional_losses_1739232) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1a_activation_layer_call_and_return_conditional_losses_1737821) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2b_activation_layer_call_and_return_conditional_losses_1751398) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5a_expand_activation_layer_call_and_return_conditional_losses_1756870) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7b_se_reduce_layer_call_and_return_conditional_losses_1763750) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6d_activation_layer_call_and_return_conditional_losses_1761467) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6a_expand_activation_layer_call_and_return_conditional_losses_1759663) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7b_expand_activation_layer_call_and_return_conditional_losses_1743011) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2c_activation_layer_call_and_return_conditional_losses_1738605) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2c_expand_activation_layer_call_and_return_conditional_losses_1738570) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6d_expand_activation_layer_call_and_return_conditional_losses_1742167) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5c_activation_layer_call_and_return_conditional_losses_1740921) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_top_activation_layer_call_and_return_conditional_losses_1764124) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7a_expand_activation_layer_call_and_return_conditional_losses_1742821) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block1b_activation_layer_call_and_return_conditional_losses_1737978) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6e_se_reduce_layer_call_and_return_conditional_losses_1742456) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5b_activation_layer_call_and_return_conditional_losses_1740703) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6e_se_reduce_layer_call_and_return_conditional_losses_1762093) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2b_se_reduce_layer_call_and_return_conditional_losses_1738423) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4d_se_reduce_layer_call_and_return_conditional_losses_1740113) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6f_se_reduce_layer_call_and_return_conditional_losses_1762661) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4c_activation_layer_call_and_return_conditional_losses_1755302) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2b_expand_activation_layer_call_and_return_conditional_losses_1751262) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4a_activation_layer_call_and_return_conditional_losses_1754213) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4c_se_reduce_layer_call_and_return_conditional_losses_1739895) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2a_activation_layer_call_and_return_conditional_losses_1750877) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2a_expand_activation_layer_call_and_return_conditional_losses_1738161) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4b_se_reduce_layer_call_and_return_conditional_losses_1754792) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4c_expand_activation_layer_call_and_return_conditional_losses_1755166) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block7a_se_reduce_layer_call_and_return_conditional_losses_1763229) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3a_se_reduce_layer_call_and_return_conditional_losses_1738860) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3c_activation_layer_call_and_return_conditional_losses_1753634) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4d_se_reduce_layer_call_and_return_conditional_losses_1755928) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4a_expand_activation_layer_call_and_return_conditional_losses_1754066) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_action_1723318) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block4b_activation_layer_call_and_return_conditional_losses_1739641) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6d_se_reduce_layer_call_and_return_conditional_losses_1742238) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6a_se_reduce_layer_call_and_return_conditional_losses_1741612) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2b_activation_layer_call_and_return_conditional_losses_1738387) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6e_expand_activation_layer_call_and_return_conditional_losses_1761899) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2c_activation_layer_call_and_return_conditional_losses_1751966) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block3a_activation_layer_call_and_return_conditional_losses_1738824) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block5a_expand_activation_layer_call_and_return_conditional_losses_1740478) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_action_1717394) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block6b_expand_activation_layer_call_and_return_conditional_losses_1760195) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_distribution_fn_1729404) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_block2b_se_reduce_layer_call_and_return_conditional_losses_1751456) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n", "WARNING:absl:Importing a function (__inference_action_1708239) with ops with unsaved custom gradients. Will likely fail if a gradient is requested.\n" ] } ], "source": [ "# @title Select your model and environment\n", "model_name = \"rt_1_x\" # @param [\"rt_1_x\", \"rt_1_400k\", \"rt_1_58k\", \"rt_1_1k\", \"octo-base\", \"octo-small\"]\n", "\n", "if \"rt_1\" in model_name:\n", " from simpler_env.policies.rt1.rt1_model import RT1Inference\n", "\n", " ckpt_path = get_rt_1_checkpoint(model_name)\n", " model = RT1Inference(saved_model_path=ckpt_path, policy_setup=policy_setup)\n", "elif \"octo\" in model_name:\n", " from simpler_env.policies.octo.octo_model import OctoInference\n", "\n", " model = OctoInference(model_type=model_name, policy_setup=policy_setup, init_rng=0)\n", "else:\n", " raise ValueError(model_name)\n" ] }, { "cell_type": "code", "execution_count": 15, "metadata": { "cellView": "form", "colab": { "base_uri": "https://localhost:8080/" }, "id": "NfXbTg_ddy5O", "outputId": "f7406200-53c8-4c97-8b9f-d55f39b3b57f" }, "outputs": [ { "output_type": "stream", "name": "stderr", "text": [ "/usr/local/lib/python3.10/dist-packages/gymnasium/core.py:311: UserWarning: \u001b[33mWARN: env.get_language_instruction to get variables from other wrappers is deprecated and will be removed in v1.0, to get this variable you can do `env.unwrapped.get_language_instruction` for environment variables or `env.get_wrapper_attr('get_language_instruction')` that will search the reminding wrappers.\u001b[0m\n", " logger.warn(\n" ] }, { "output_type": "stream", "name": "stdout", "text": [ "pick coke can\n" ] }, { "output_type": "stream", "name": "stderr", "text": [ "/usr/local/lib/python3.10/dist-packages/gymnasium/core.py:311: UserWarning: \u001b[33mWARN: env.robot_uid to get variables from other wrappers is deprecated and will be removed in v1.0, to get this variable you can do `env.unwrapped.robot_uid` for environment variables or `env.get_wrapper_attr('robot_uid')` that will search the reminding wrappers.\u001b[0m\n", " logger.warn(\n" ] }, { "output_type": "stream", "name": "stdout", "text": [ "0 {'elapsed_steps': 1, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "1 {'elapsed_steps': 2, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "2 {'elapsed_steps': 3, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "3 {'elapsed_steps': 4, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "4 {'elapsed_steps': 5, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "5 {'elapsed_steps': 6, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "6 {'elapsed_steps': 7, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "7 {'elapsed_steps': 8, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "8 {'elapsed_steps': 9, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "9 {'elapsed_steps': 10, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "10 {'elapsed_steps': 11, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "11 {'elapsed_steps': 12, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "12 {'elapsed_steps': 13, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "13 {'elapsed_steps': 14, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "14 {'elapsed_steps': 15, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "15 {'elapsed_steps': 16, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "16 {'elapsed_steps': 17, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "17 {'elapsed_steps': 18, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "18 {'elapsed_steps': 19, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "19 {'elapsed_steps': 20, 'is_grasped': False, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', False)])}\n", "20 {'elapsed_steps': 21, 'is_grasped': True, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', True)])}\n", "21 {'elapsed_steps': 22, 'is_grasped': True, 'consecutive_grasp': False, 'lifted_object': False, 'lifted_object_significantly': False, 'success': False, 'episode_stats': OrderedDict([('n_lift_significant', 0), ('consec_grasp', False), ('grasped', True)])}\n", "22 {'elapsed_steps': 23, 'is_grasped': True, 'consecutive_grasp': False, 'lifted_object': True, 'lifted_object_significantly': True, 'success': True, 'episode_stats': OrderedDict([('n_lift_significant', 1), ('consec_grasp', False), ('grasped', True)])}\n", "23 {'elapsed_steps': 24, 'is_grasped': True, 'consecutive_grasp': False, 'lifted_object': True, 'lifted_object_significantly': True, 'success': True, 'episode_stats': OrderedDict([('n_lift_significant', 2), ('consec_grasp', False), ('grasped', True)])}\n", "24 {'elapsed_steps': 25, 'is_grasped': True, 'consecutive_grasp': True, 'lifted_object': True, 'lifted_object_significantly': True, 'success': True, 'episode_stats': OrderedDict([('n_lift_significant', 3), ('consec_grasp', True), ('grasped', True)])}\n", "25 {'elapsed_steps': 26, 'is_grasped': True, 'consecutive_grasp': True, 'lifted_object': True, 'lifted_object_significantly': True, 'success': True, 'episode_stats': OrderedDict([('n_lift_significant', 4), ('consec_grasp', True), ('grasped', True)])}\n", "Episode success: True\n" ] } ], "source": [ "#@title Run inference\n", "\n", "obs, reset_info = env.reset()\n", "instruction = env.get_language_instruction()\n", "model.reset(instruction)\n", "print(instruction)\n", "\n", "image = get_image_from_maniskill2_obs_dict(env, obs) # np.ndarray of shape (H, W, 3), uint8\n", "images = [image]\n", "predicted_terminated, success, truncated = False, False, False\n", "timestep = 0\n", "while not (predicted_terminated or truncated):\n", " # step the model; \"raw_action\" is raw model action output; \"action\" is the processed action to be sent into maniskill env\n", " raw_action, action = model.step(image)\n", " predicted_terminated = bool(action[\"terminate_episode\"][0] > 0)\n", " obs, reward, success, truncated, info = env.step(\n", " np.concatenate([action[\"world_vector\"], action[\"rot_axangle\"], action[\"gripper\"]])\n", " )\n", " print(timestep, info)\n", " # update image observation\n", " image = get_image_from_maniskill2_obs_dict(env, obs)\n", " images.append(image)\n", " timestep += 1\n", "\n", "episode_stats = info.get(\"episode_stats\", {})\n", "print(f\"Episode success: {success}\")" ] }, { "cell_type": "code", "execution_count": 16, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 553 }, "id": "ZiiQ0Fnqy2Wb", "outputId": "bd5e489d-6c2c-45c2-eab2-fd114c4f89ee" }, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "google_robot_pick_coke_can rt_1_x\n" ] }, { "output_type": "display_data", "data": { "text/plain": [ "" ], "text/html": [ "
" ] }, "metadata": {} } ], "source": [ "print(task_name, model_name)\n", "mediapy.show_video(images, fps=10)" ] }, { "cell_type": "markdown", "metadata": { "id": "v69YCHasyz3O" }, "source": [ "## Gallery" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 535 }, "id": "pe59bPnNyka6", "outputId": "46f14748-9c0b-4017-9427-2750f423e7a7" }, "outputs": [ { "data": { "text/html": [ "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# @markdown RT-1-X close drawer\n", "print(task_name, model_name)\n", "mediapy.show_video(images, fps=10)\n", "# Note: we turned off the denoiser as the colab kernel will crash if it's turned on\n", "# To use the denoiser, please git clone our SIMPLER environments\n", "# and perform evaluations locally." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 521 }, "id": "b8pDRz-dzUJZ", "outputId": "0515dec6-5ccf-4811-b0ae-67710965f245" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "widowx_put_eggplant_in_basket octo-base\n" ] }, { "data": { "text/html": [ "
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# @markdown Octo-base widowx_put_eggplant_in_basket\n", "print(task_name, model_name)\n", "mediapy.show_video(images, fps=10)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [ "TZJMaVVObgRt" ], "gpuType": "T4", "provenance": [], "include_colab_link": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }