inference_pretrained.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "534daf7f-4b6b-4357-9a38-9117f72ce9b4",
   "metadata": {},
   "source": [
    "# Step 1: Minimal CrossFormer Inference Example\n",
    "\n",
    "This Colab demonstrates how to load a pre-trained / finetuned CrossFormer checkpoint, run inference for a single-arm and bimanual manipulation system, and compare the outputs to the true actions.\n",
    "\n",
    "First, let's start with a minimal example!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca93d021",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run this block if you're using Colab\n",
    "\n",
    "# Download repo\n",
    "!git clone https://github.com/rail-berkeley/crossformer.git\n",
    "%cd crossformer\n",
    "# Install repo\n",
    "!pip3 install -e .\n",
    "!pip3 install -r requirements.txt\n",
    "!pip3 install --upgrade \"jax[cuda11_pip]==0.4.20\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83d34283",
   "metadata": {},
   "outputs": [],
   "source": [
    "from crossformer.model.crossformer_model import CrossFormerModel\n",
    "\n",
    "model = CrossFormerModel.load_pretrained(\"hf://rail-berkeley/crossformer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e669650f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll demonstrate how to create an observation and task dictionary for a bimanual task. \n",
    "# Then we'll use them to sample an action from the model.\n",
    "\n",
    "import jax\n",
    "import numpy as np\n",
    "\n",
    "# create a random image\n",
    "img = np.random.randint(0, 255, size=(224, 224, 3))\n",
    "# add batch and observation history dimension (CrossFormer accepts a history of up to 5 time-steps)\n",
    "img = img[None, None]\n",
    "# our bimanual training data has an overhead view and two wrist views\n",
    "observation = {\n",
    "    \"image_high\": img,\n",
    "    \"image_left_wrist\": img,\n",
    "    \"image_right_wrist\": img,\n",
    "    \"timestep_pad_mask\": np.array([[True]]),\n",
    "}\n",
    "# create a task dictionary for a language task\n",
    "task = model.create_tasks(texts=[\"uncap the pen\"])\n",
    "# we need to specify the bimanual head here\n",
    "action = model.sample_actions(observation, task, head_name=\"bimanual\", rng=jax.random.PRNGKey(0))\n",
    "print(action)  # [batch, action_chunk, action_dim]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2be0d1f",
   "metadata": {},
   "source": [
    "# Step 2: Run Inference on Full Trajectories\n",
    "\n",
    "That was easy! Now let's try to run inference across a whole single-arm trajectory and visualize the results!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a51eb166",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install mediapy for visualization\n",
    "!pip install mediapy\n",
    "!pip install opencv-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b0f7fd1-5b43-480f-b00f-766248d7f9af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import jax\n",
    "import tensorflow_datasets as tfds\n",
    "import tqdm\n",
    "import mediapy\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b79053f4-316f-4d2d-81bd-e6e04cfa81bf",
   "metadata": {},
   "source": [
    "## Load Model Checkpoint\n",
    "First, we will load the pre-trained checkpoint using the `load_pretrained()` function. You can specify the path to a checkpoint directory or a HuggingFace path.\n",
    "\n",
    "Below, we are loading directly from HuggingFace.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42c04953-869d-48a8-a2df-e601324e97e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from crossformer.model.crossformer_model import CrossFormerModel\n",
    "\n",
    "model = CrossFormerModel.load_pretrained(\"hf://rail-berkeley/crossformer\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c298ac8f-da06-41d5-a4a5-145c3080231e",
   "metadata": {},
   "source": [
    "## Load Data\n",
    "Next, we will load a trajectory from the Bridge dataset for testing the model. We will use the publicly available copy in the Open X-Embodiment dataset bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "392bd127",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create RLDS dataset builder\n",
    "builder = tfds.builder_from_directory(\n",
    "    builder_dir=\"gs://gresearch/robotics/bridge/0.1.0/\"\n",
    ")\n",
    "ds = builder.as_dataset(split=\"train[:1]\")\n",
    "\n",
    "# sample episode and resize to 224x224 (default third-person cam resolution)\n",
    "episode = next(iter(ds))\n",
    "steps = list(episode[\"steps\"])\n",
    "images = [\n",
    "    cv2.resize(np.array(step[\"observation\"][\"image\"]), (224, 224)) for step in steps\n",
    "]\n",
    "\n",
    "# extract goal image and language instruction\n",
    "goal_image = images[-1]\n",
    "language_instruction = (\n",
    "    steps[0][\"observation\"][\"natural_language_instruction\"].numpy().decode()\n",
    ")\n",
    "\n",
    "# visualize episode\n",
    "print(f\"Instruction: {language_instruction}\")\n",
    "mediapy.show_video(images, fps=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b37ffca5",
   "metadata": {},
   "source": [
    "## Run Inference\n",
    "\n",
    "Next, we will run inference over the images in the episode using the loaded model. \n",
    "Below we demonstrate setups for both goal-conditioned and language-conditioned training.\n",
    "Note that we need to feed inputs of the correct temporal window size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ad64434",
   "metadata": {},
   "outputs": [],
   "source": [
    "WINDOW_SIZE = 5\n",
    "\n",
    "# create task dictionary\n",
    "task = model.create_tasks(\n",
    "    goals={\"image_primary\": goal_image[None]}\n",
    ")  # for goal-conditioned\n",
    "task = model.create_tasks(texts=[language_instruction])  # for language conditioned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74d6b20f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run inference loop, the model only uses 3rd person image observations for bridge\n",
    "\n",
    "# collect predicted and true actions\n",
    "pred_actions, true_actions = [], []\n",
    "for step in tqdm.trange(len(images) - (WINDOW_SIZE - 1)):\n",
    "    input_images = np.stack(images[step : step + WINDOW_SIZE])[None]\n",
    "    observation = {\n",
    "        \"image_primary\": input_images,\n",
    "        \"timestep_pad_mask\": np.full((1, input_images.shape[1]), True, dtype=bool),\n",
    "    }\n",
    "\n",
    "    # we need to pass in the dataset statistics to unnormalize the actions\n",
    "    actions = model.sample_actions(\n",
    "        observation,\n",
    "        task,\n",
    "        head_name=\"single_arm\",\n",
    "        unnormalization_statistics=model.dataset_statistics[\"bridge_dataset\"][\"action\"],\n",
    "        rng=jax.random.PRNGKey(0),\n",
    "    )\n",
    "    actions = actions[0]  # remove batch\n",
    "\n",
    "    pred_actions.append(actions)\n",
    "    final_window_step = step + WINDOW_SIZE - 1\n",
    "    true_actions.append(\n",
    "        np.concatenate(\n",
    "            (\n",
    "                steps[final_window_step][\"action\"][\"world_vector\"],\n",
    "                steps[final_window_step][\"action\"][\"rotation_delta\"],\n",
    "                np.array(steps[final_window_step][\"action\"][\"open_gripper\"]).astype(\n",
    "                    np.float32\n",
    "                )[None],\n",
    "            ),\n",
    "            axis=-1,\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12a5e3f7",
   "metadata": {},
   "source": [
    "## Visualize predictions and ground-truth actions\n",
    "\n",
    "Finally, we will visualize the predicted actions in comparison to the groundtruth actions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a79775d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "ACTION_DIM_LABELS = ['x', 'y', 'z', 'yaw', 'pitch', 'roll', 'grasp']\n",
    "\n",
    "# build image strip to show above actions\n",
    "img_strip = np.concatenate(np.array(images[::3]), axis=1)\n",
    "\n",
    "# set up plt figure\n",
    "figure_layout = [\n",
    "    ['image'] * len(ACTION_DIM_LABELS),\n",
    "    ACTION_DIM_LABELS\n",
    "]\n",
    "plt.rcParams.update({'font.size': 12})\n",
    "fig, axs = plt.subplot_mosaic(figure_layout)\n",
    "fig.set_size_inches([45, 10])\n",
    "\n",
    "# plot actions\n",
    "pred_actions = np.array(pred_actions).squeeze()\n",
    "true_actions = np.array(true_actions).squeeze()\n",
    "for action_dim, action_label in enumerate(ACTION_DIM_LABELS):\n",
    "  # actions have batch, horizon, dim, in this example we just take the first action for simplicity\n",
    "  axs[action_label].plot(pred_actions[:, 0, action_dim], label='predicted action')\n",
    "  axs[action_label].plot(true_actions[:, action_dim], label='ground truth')\n",
    "  axs[action_label].set_title(action_label)\n",
    "  axs[action_label].set_xlabel('Time in one episode')\n",
    "\n",
    "axs['image'].imshow(img_strip)\n",
    "axs['image'].set_xlabel('Time in one episode (subsampled)')\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}