static/notebooks/PR2LExample.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lS9FLp_n-OQV",
        "outputId": "5f1b8adf-dd8f-4cc3-9cb2-3a0b7441c0e5"
      },
      "outputs": [],
      "source": [
        "# Install transformers to enable InstructBLIP\n",
        "# Upgrade torchvision to enable Imagenette\n",
        "# !pip install transformers\n",
        "# !pip uninstall torchvision torch -y\n",
        "# !pip install torchvision==0.17.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YuJhUI2j973f"
      },
      "outputs": [],
      "source": [
        "# Imports\n",
        "import torch\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import gym\n",
        "\n",
        "from transformers import (\n",
        "    InstructBlipProcessor,\n",
        "    InstructBlipForConditionalGeneration,\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0C2klhVqK6Bx"
      },
      "source": [
        "First, let's set up a simple Gym environment. In this case, we formulate a classification problem as an MDP. At each step, the agent gets an observation (an image of size 320 x 320 from the Imagenette dataset) and has to pick one of 10 actions (corresponding to the image's class). If the agent picks correctly, it gets a reward of +1, else it gets nothing.\n",
        "\n",
        "While this is definitely a toy example, we expect PR2L to help with this task, since if the VLM is able to correctly identify the object in the image, the corresponding representations should easily be linked to the correct class by a trained policy network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xd-jbtaR-WCd",
        "outputId": "1507a3c3-7b99-4aec-f372-0473f0c6993d"
      },
      "outputs": [],
      "source": [
        "class ClassifierEnv(gym.Env):\n",
        "  \"\"\"\n",
        "  Toy contextual bandits classification task to demonstrate PR2L.\n",
        "\n",
        "  The agent receives an image from ImageNette and has to take one of 10 actions\n",
        "  corresponding to the class of the image. If correct, the agent receives +1. Else,\n",
        "  it receives 0 rewards.\n",
        "\n",
        "  Args:\n",
        "    num_imgs: int number of images from ImageNette to load. Default 1000.\n",
        "    download: bool whether or not to download the dataset. Only needs to be called\n",
        "      with True once. Default True\n",
        "    render_mode: Unused.\n",
        "  \"\"\"\n",
        "  def __init__(self, num_imgs=1000, download=True, render_mode=None, seed=0):\n",
        "    self.num_imgs = num_imgs\n",
        "\n",
        "    self.transform = transforms.Compose([transforms.CenterCrop(320)])\n",
        "\n",
        "    self.trainset = torchvision.datasets.Imagenette(root='./data', split=\"val\", size=\"320px\",\n",
        "                                                download=download, transform=self.transform)\n",
        "\n",
        "    self.classes = self.trainset.classes\n",
        "\n",
        "    self.images = []\n",
        "    self.labels = []\n",
        "\n",
        "    np.random.seed(seed)\n",
        "    for i in np.random.permutation(len(self.trainset)):\n",
        "      img, label = self.trainset.__getitem__(i)\n",
        "      self.images.append(np.array(img))\n",
        "      self.labels.append(label)\n",
        "      if len(self.labels) >= num_imgs:\n",
        "        break\n",
        "\n",
        "    self.current_idx = 0\n",
        "\n",
        "    self.observation_space = gym.spaces.Box(low=0, high=255, shape=[320, 320, 3])\n",
        "    self.action_space = gym.spaces.Discrete(len(self.classes))\n",
        "\n",
        "    self.render_mode = render_mode\n",
        "\n",
        "  def reset(self, seed=0, options=None):\n",
        "    self.current_idx = np.random.choice(len(self.images))\n",
        "    return self.images[self.current_idx], {}\n",
        "\n",
        "  def step(self, action):\n",
        "    correct_action = self.labels[self.current_idx]\n",
        "\n",
        "    reward = 1 if action == correct_action else 0\n",
        "\n",
        "    info = {}\n",
        "    done = True\n",
        "    trunc = False\n",
        "\n",
        "    return self.images[0] * 0, reward, done, trunc, info\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next, let's write a wrapper that embeds images from the base environment with the VLM (using a provided prompt). In this case, the VLM is effectively part of the environment, with no gradients flowing through it. It simply acts as an encoder of images to yield promptable state representations. \n",
        "\n",
        "In principle, one could have the VLM be part of the policy and be updated with RL as well (akin to RT-2), but this is computationally expensive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iAKFhnb9N277"
      },
      "outputs": [],
      "source": [
        "class VlmImgFeatureWrapper(gym.ObservationWrapper):\n",
        "    \"\"\"\n",
        "    Wrapper for PR2L\n",
        "\n",
        "    Each image in the wrapped environment will be parsed via\n",
        "    the VLM with a given input prompt, thereby converting it to a sequence\n",
        "    of token embeddings.\n",
        "\n",
        "    Specifically, the wrapped environment's observation space replaces \"rgb\" with\n",
        "    \"seq\" and \"seq_mask\". The former is a sequence of token embeddings of shape\n",
        "    (num tokens, hidden_dim), where num tokens depends on some of the below parameters\n",
        "    and hidden_dim varies with the chosen VLM. seq_mask then indicates which token embeddings\n",
        "    correspond to padding tokens (if the generate_kwargs permits those). These two arrays\n",
        "    are directly usable by PyTorch transformer layers (as the encoder input and mask).\n",
        "\n",
        "    Some of the below args depend on the VLM being used.\n",
        "    InstructBLIP-vicuna7b (recommended): hidden_dim = 4096, skip_first = True\n",
        "    InstructBLIP-vicuna13b: hidden_dim = 5120, skip_first = True\n",
        "    InstructBLIP-flan-t5-xl: hidden_dim = 2048, skip_first = False\n",
        "\n",
        "    Args:\n",
        "    env: gym Env to be wrapped.\n",
        "    vlm_model: transformers VLM being used to embed images\n",
        "    processor: transformers processor for the VLM\n",
        "    hidden_dim: int dimensionality of token embeddings for the VLM\n",
        "    prompt: str prompt given to the VLM with the image\n",
        "    last_n_layers: int number of layers of the transformer-based VLM whose outputed token\n",
        "        embeddings are included in the extracted promptable representation. Default 1\n",
        "        (use the token embeddings outputted by the final layer only)\n",
        "    use_encoder_embeds: bool, whether or not to use the token embeddings corresponding\n",
        "        to the prompt or input image. Default false (just use the embeddings of generated text)\n",
        "    skip_first: bool, whether or not to skip the first element of the hidden states object.\n",
        "        Depends on the VLM used. Default true.\n",
        "    move_to_cpu: bool, whether to move generated embeddings to CPU (as numpy array) or keep\n",
        "        it on same device as the VLM (as torch tensor). Default true.\n",
        "    verbose: bool, whether or not to print out generated text. Default false.\n",
        "    generate_kwargs: dict of generate keyword args used by VLM (see Huggingface documentation)\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        env,\n",
        "        vlm_model,\n",
        "        processor,\n",
        "        hidden_dim: int,\n",
        "        prompt: str,\n",
        "        img_shape,\n",
        "        last_n_layers=1,\n",
        "        use_encoder_embeds=False,\n",
        "        skip_first=True,\n",
        "        generate_kwargs=dict(\n",
        "            max_new_tokens=8,\n",
        "            min_new_tokens=8,\n",
        "            output_hidden_states=True,\n",
        "            return_dict_in_generate=True,\n",
        "            do_sample=False,\n",
        "        ),\n",
        "    ):\n",
        "        super().__init__(env)\n",
        "\n",
        "        self.img_shape = img_shape\n",
        "        self.vlm_model = vlm_model\n",
        "        self.device = self.vlm_model.device\n",
        "        self.processor = processor\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.prompt = prompt\n",
        "        self.last_n_layers = last_n_layers\n",
        "        self.use_encoder_embeds = use_encoder_embeds\n",
        "        self.generate_kwargs = generate_kwargs\n",
        "        self.skip_first = skip_first\n",
        "\n",
        "\n",
        "        self.observation_space = self._make_space(self.observation_space)\n",
        "\n",
        "    def observation(self, obs):\n",
        "        # modify obs\n",
        "        return self._process_obs(obs)\n",
        "\n",
        "    def _make_space(self, obs_space):\n",
        "        map = {}\n",
        "        new_shape = self._get_new_shape()\n",
        "        self.max_embs = new_shape[0]\n",
        "        print(\"Initializing VLM embedding observation space with...\")\n",
        "        print(f\"max_new_tokens = {self.generate_kwargs['max_new_tokens']}\")\n",
        "        print(f\"last_n_layers = {self.last_n_layers}\")\n",
        "        print(f\"Thus, setting final embedding observation shape to {new_shape}\")\n",
        "\n",
        "        new_v = gym.spaces.Box(low=-np.inf, high=np.inf, shape=new_shape)\n",
        "        map[\"seq\"] = new_v\n",
        "        map[\"seq_mask\"] = gym.spaces.Box(\n",
        "            low=np.ones(self.max_embs)*False, high=np.ones(self.max_embs)*True, dtype=bool\n",
        "        )\n",
        "        return gym.spaces.Dict(map)\n",
        "\n",
        "    def _process_obs(self, obs):\n",
        "        \"\"\"\n",
        "        Convert from an image observation to a dictionary with keys:\n",
        "            seq: Sequence of VLM representations corresponding to image of size\n",
        "                (max sequence length, token embed dim)\n",
        "            seq_mask: Mask of padding, vector of bools of length (max sequence length)\n",
        "        \"\"\"\n",
        "        map = {}\n",
        "        img = obs.copy()\n",
        "        if len(img.shape) < 4:\n",
        "            img = img.reshape([1] * (4 - len(img.shape)) + [*img.shape])\n",
        "        map[\"seq\"], map[\"seq_mask\"] = self._generate_embeds(img)\n",
        "        return map\n",
        "\n",
        "    def _get_new_shape(self):\n",
        "        \"\"\"\n",
        "        Gets the shape of the VLM representation sequence, for the purpose of creating\n",
        "        a suitable observation space\n",
        "        \"\"\"\n",
        "        seq, mask = self._generate_embeds(np.zeros(self.img_shape))\n",
        "        return seq.shape\n",
        "\n",
        "    def _generate_embeds(self, img):\n",
        "        \"\"\"\n",
        "        Pass prompt and image through VLM to yield hidden states, packaging it into\n",
        "        dictionary observation\n",
        "        \"\"\"\n",
        "        inputs = self.processor(images=img, text=self.prompt, return_tensors=\"pt\").to(\n",
        "            self.device\n",
        "        )\n",
        "        inputs[\"pixel_values\"] = inputs[\"pixel_values\"].to(\n",
        "            self.device, self.vlm_model.dtype\n",
        "        )\n",
        "        generated_ids = self.vlm_model.generate(**inputs, **self.generate_kwargs)\n",
        "        if \"hidden_states\" in generated_ids.keys():\n",
        "            hs = generated_ids[\"hidden_states\"]\n",
        "        elif \"decoder_hidden_states\" in generated_ids.keys():\n",
        "            assert not self.skip_first\n",
        "            hs = {\n",
        "                \"encoder\": generated_ids[\"encoder_hidden_states\"],\n",
        "                \"decoder\": generated_ids[\"decoder_hidden_states\"]\n",
        "                }\n",
        "        else:\n",
        "            raise NotImplementedError\n",
        "\n",
        "        return self._get_embeds(hs)\n",
        "\n",
        "\n",
        "    def _get_embeds(self, hs):\n",
        "        \"\"\"\n",
        "        Produces a single embedding tensor of shape:\n",
        "        [seq len, num hidden states, hidden state dims]\n",
        "\n",
        "        Args:\n",
        "            hs: tuple of tuple of tensors. Outer tuple has shape\n",
        "                # generated tokens (+ 1 if skip_first). Inner tuple\n",
        "                has shape (number of layers of self.vlm_model + 1).\n",
        "                Tensors in all layers other than\n",
        "            skip_first: bool (default True) on whether the first token\n",
        "                of hs should be skipped (ie if its the hidden states\n",
        "                for the prompt -- set to true for InstructBLIP/BLIP2)\n",
        "        Returns:\n",
        "            tuple of sequence of token embeds from VLM and corresponding\n",
        "                padding mask\n",
        "        \"\"\"\n",
        "        tokenwise_emb, mask = [], []\n",
        "\n",
        "        if not self.skip_first:\n",
        "            assert type(hs) is dict\n",
        "            # Used for T5-Flan InstructBLIP versions (as in this demo notebook)\n",
        "            embs = []\n",
        "            for i in range(len(hs[\"decoder\"])):\n",
        "                # [1, self.last_n_layers, token dim]\n",
        "                last_n_reps = torch.cat(hs[\"decoder\"][i][-self.last_n_layers:], dim = 1)\n",
        "                embs.append(last_n_reps)\n",
        "\n",
        "            # Create padding mask (True represents corresponding token is padding)\n",
        "            padding = [False] * len(embs) * self.last_n_layers\n",
        "            num_pad_tokens = self.generate_kwargs[\"max_new_tokens\"] - len(embs)\n",
        "            padding += [True] * self.last_n_layers * num_pad_tokens\n",
        "\n",
        "            if num_pad_tokens > 0:\n",
        "                embs += [torch.zeros(1, self.last_n_layers, self.hidden_dim)] * num_pad_tokens\n",
        "\n",
        "            # [number generated tokens, self.last_n_layers, token dim]\n",
        "            embs = torch.cat(embs, dim=0)\n",
        "\n",
        "            if self.use_encoder_embeds: # Whether or not to include representations from prompt/image\n",
        "                # [self.last_n_layers, num encoder tokens, token dim]\n",
        "                enc_embs = torch.cat(hs[\"encoder\"][-self.last_n_layers:], dim=0)\n",
        "                enc_embs = enc_embs.permute(1, 0, 2)\n",
        "                padding = [False] * len(enc_embs) * self.last_n_layers + padding\n",
        "\n",
        "                embs = torch.cat([enc_embs, embs], dim=0)\n",
        "\n",
        "            # [total tokens * self.last_n_layers, token dim]\n",
        "            embs = embs.reshape(-1, self.hidden_dim)\n",
        "\n",
        "            assert len(embs) == len(padding)\n",
        "\n",
        "            return embs.detach().cpu().numpy(), np.array(padding)\n",
        "\n",
        "        else:\n",
        "            # Used for Vicuna InstructBLIP versions\n",
        "            if self.use_encoder_embeds:\n",
        "                # shape: (# layers -> last n layers, # enc tokens, # hidden dims)\n",
        "                encoder_hs = torch.stack(hs[0], dim=0)[-self.last_n_layers :].detach()\n",
        "                # shape: (# enc tokens, # layers, # hidden dims)\n",
        "                encoder_hs = torch.transpose(encoder_hs, 0, 1)\n",
        "                # shape: (# enc tokens * # layers, # hidden dims)\n",
        "                encoder_hs = encoder_hs.reshape(-1, encoder_hs.shape[-1])\n",
        "                tokenwise_emb.append(encoder_hs.cpu().numpy())\n",
        "                mask += [False] * len(encoder_hs)\n",
        "            tokens = hs[1:]\n",
        "\n",
        "            for token in tokens:\n",
        "                for emb in token[-self.last_n_layers :]:\n",
        "                    final_token_emb = emb.detach().reshape(1, -1)\n",
        "                    if self.move_to_cpu:\n",
        "                        final_token_emb = final_token_emb.cpu().numpy()\n",
        "                    tokenwise_emb.append(final_token_emb)\n",
        "                    mask.append(False)\n",
        "\n",
        "            while len(mask) < self.max_embs:\n",
        "                emb = np.zeros([1, self.hidden_dim])\n",
        "                tokenwise_emb.append(emb)\n",
        "                mask.append(True)\n",
        "            return np.concatenate(tokenwise_emb), np.array(mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 249
        },
        "id": "epNZsLlb316H",
        "outputId": "4c2a773f-50c2-4067-c98d-5a560ebe6afb"
      },
      "outputs": [],
      "source": [
        "# Initialize the VLM\n",
        "VLM_DEVICE = \"cuda:7\"\n",
        "VLM_DTYPE = torch.float16\n",
        "\n",
        "vlm_model = InstructBlipForConditionalGeneration.from_pretrained(\n",
        "            \"Salesforce/instructblip-flan-t5-xl\", torch_dtype=VLM_DTYPE\n",
        "        ).to(VLM_DEVICE)\n",
        "processor = InstructBlipProcessor.from_pretrained(\n",
        "    \"Salesforce/instructblip-flan-t5-xl\"\n",
        ")\n",
        "vlm_dim = 2048\n",
        "skip_first = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lbxo7NBz52z7"
      },
      "outputs": [],
      "source": [
        "# Create env\n",
        "env = ClassifierEnv(download=False)\n",
        "\n",
        "# Wrap env with PR2L wrapper\n",
        "env = VlmImgFeatureWrapper(\n",
        "    env, \n",
        "    vlm_model=vlm_model, \n",
        "    processor=processor, \n",
        "    hidden_dim=vlm_dim, \n",
        "    img_shape=[320, 320, 3], \n",
        "    prompt=\"What is in this image?\", \n",
        "    # Answering this prompt correctly should yield good \n",
        "    # representations for getting linked to the proper category\n",
        "    last_n_layers=2,\n",
        "    use_encoder_embeds=True,\n",
        "    skip_first=skip_first\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now, we create a custom policy network that takes in the VLM representation sequence and processes it with a learned Transformer layer. We opt to use an encoder-decoder architecture that summarizes the full sequence in a single \"CLS\" token's embedding, which is computed by attending to all other tokens in the sequence. This strategy allows the full sequence to be condensed into a single fixed-size embedding, which can be passed into a downstream model head to produce actions (akin to how the \"CLS\" token in the BERT LLM is used for classification in natural language processing tasks).\n",
        "\n",
        "The downstream model is initialized with the default Stable Baselines 3 Actor-Critic policy, which expects fixed-size inputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from stable_baselines3.common.policies import ActorCriticPolicy\n",
        "from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union\n",
        "from stable_baselines3.common.type_aliases import Schedule\n",
        "from stable_baselines3.common.torch_layers import (\n",
        "    BaseFeaturesExtractor,\n",
        "    CombinedExtractor,\n",
        ")\n",
        "import gymnasium\n",
        "\n",
        "class TfInputActorCriticPolicy(ActorCriticPolicy):\n",
        "    \"\"\"\n",
        "    Create a policy with a transformer layer for handling sequential VLM\n",
        "    representations.\n",
        "\n",
        "    Most arguments are for the Stable Baselines 3 ActorCriticPolicy, since \n",
        "    we initialize said policy to take in the token embedding of the \n",
        "    transformer's CLS token. Below are the auxiliary args added for the \n",
        "    transformer layer.\n",
        "\n",
        "    Args:\n",
        "        tf_input_key: str representing key of the VLM representation sequence\n",
        "            in the observation dictionary, default seq\n",
        "        tf_mask_key: str representing key of the padding mask in the observation\n",
        "            dictionary, default seq_mask (can be None if no mask is needed)\n",
        "        raw_seq_dim: int dimensionality of individual VLM token embeddings\n",
        "        tf_to_fc_dim: int dimensionality that CLS token embedding is projected to\n",
        "            before being given to MLP part of policy\n",
        "        tf_kwargs: dict of additional transformer layer parameters. See\n",
        "            https://pytorch.org/docs/stable/generated/torch.nn.Transformer.html\n",
        "    \"\"\"\n",
        "    def __init__(\n",
        "        self,\n",
        "        observation_space: gym.spaces.Dict,\n",
        "        action_space: gym.spaces.Space,\n",
        "        lr_schedule = 1e-4,\n",
        "        tf_input_key: str = \"seq\",\n",
        "        tf_mask_key=\"seq_mask\",\n",
        "        net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,\n",
        "        activation_fn: Type[torch.nn.Module] = torch.nn.Tanh,\n",
        "        ortho_init: bool = True,\n",
        "        use_sde: bool = False,\n",
        "        log_std_init: float = 0.0,\n",
        "        full_std: bool = True,\n",
        "        use_expln: bool = False,\n",
        "        squash_output: bool = False,\n",
        "        features_extractor_class: Type[BaseFeaturesExtractor] = CombinedExtractor,\n",
        "        features_extractor_kwargs: Optional[Dict[str, Any]] = None,\n",
        "        share_features_extractor: bool = True,\n",
        "        normalize_images: bool = True,\n",
        "        optimizer_class: Type[torch.optim.Optimizer] = torch.optim.Adam,\n",
        "        optimizer_kwargs: Optional[Dict[str, Any]] = None,\n",
        "        raw_seq_dim: int = 768,\n",
        "        tf_to_fc_dim: int = 512,\n",
        "        tf_kwargs={\n",
        "            \"d_model\": 768,\n",
        "            \"nhead\": 2,\n",
        "            \"num_encoder_layers\": 1,\n",
        "            \"num_decoder_layers\": 1,\n",
        "            \"dim_feedforward\": 512,\n",
        "            \"batch_first\": True,\n",
        "        },\n",
        "    ):\n",
        "        self.tf_to_fc_dim = tf_to_fc_dim\n",
        "        self.raw_seq_dim = raw_seq_dim\n",
        "        self.tf_input_key = tf_input_key\n",
        "        self.tf_mask_key = tf_mask_key\n",
        "        self.tf_kwargs = tf_kwargs\n",
        "\n",
        "        # Compute observation space after passing VLM representations through\n",
        "        # transformer layer\n",
        "        observation_space = self._make_obs_space(observation_space)\n",
        "\n",
        "        # Initialize standard MLP ActorCriticPolicy that takes in CLS token\n",
        "        # embed outputted by transformer layer.\n",
        "        super().__init__(\n",
        "            observation_space,\n",
        "            action_space,\n",
        "            lr_schedule,\n",
        "            net_arch,\n",
        "            activation_fn,\n",
        "            ortho_init,\n",
        "            use_sde,\n",
        "            log_std_init,\n",
        "            full_std,\n",
        "            use_expln,\n",
        "            squash_output,\n",
        "            features_extractor_class,\n",
        "            features_extractor_kwargs,\n",
        "            share_features_extractor,\n",
        "            normalize_images,\n",
        "            optimizer_class,\n",
        "            optimizer_kwargs,\n",
        "        )\n",
        "        \n",
        "        # Initialize layer lowering dimensionality of VLM token embeds\n",
        "        self.fc1 = torch.nn.Linear(raw_seq_dim, tf_kwargs[\"d_model\"])\n",
        "\n",
        "        # Initialize transformer layer\n",
        "        self.tf = torch.nn.Transformer(**tf_kwargs)\n",
        "\n",
        "        # Initialize layer changing dimensionality of CLS token embed before\n",
        "        # passing through MLP policy\n",
        "        self.fc2 = torch.nn.Linear(tf_kwargs[\"d_model\"], self.tf_to_fc_dim)\n",
        "        \n",
        "        # Initialize learned CLS token embed\n",
        "        self.decode_emb = torch.nn.Embedding(1, tf_kwargs[\"d_model\"])\n",
        "        self.moved_emb = False\n",
        "\n",
        "        # Build again to add the transformer params to the optimizer\n",
        "        self._build(lr_schedule)\n",
        "\n",
        "    def _make_obs_space(self, observation_space: gym.spaces.Dict):\n",
        "        new_obs_space = {}\n",
        "        for k, v in observation_space.items():\n",
        "            if k == self.tf_input_key:\n",
        "                new_obs_space[k] = gymnasium.spaces.Box(\n",
        "                    low=-np.inf, high=np.inf, shape=[self.tf_to_fc_dim]\n",
        "                )\n",
        "            elif k == self.tf_mask_key:\n",
        "                continue\n",
        "            else:\n",
        "                new_obs_space[k] = v\n",
        "\n",
        "        return gymnasium.spaces.Dict(new_obs_space)  # gym.spaces.Dict(new_obs_space)\n",
        "\n",
        "    def _process_seq(self, tf_input, tf_mask=None):\n",
        "        # Project tf_input token size from raw_seq_dim -> d_model\n",
        "        tf_input = self.fc1(tf_input)\n",
        "\n",
        "        # Determine number of CLS tokens necessary\n",
        "        if len(tf_input.shape) == 3:\n",
        "            batch_size = tf_input.shape[0]\n",
        "            tgt_emb_shape = [batch_size, 1]\n",
        "        else:\n",
        "            tgt_emb_shape = [1]\n",
        "\n",
        "        # Move CLS token to GPU, if not already there\n",
        "        if not self.moved_emb:\n",
        "            self.decode_emb.to(self.device)\n",
        "            self.moved_emb = True\n",
        "\n",
        "        # Get CLS tokens for each sequence in batch\n",
        "        tgt_emb = self.decode_emb(torch.zeros(tgt_emb_shape, dtype=int).to(self.device))\n",
        "\n",
        "        # Pass input sequence through policy transformer layer encoder and CLS through decoder\n",
        "        # Use mask if provided.\n",
        "        # Computes the CLS token's embeddings, as computed by attending to the VLM's promptable reps.\n",
        "        # Shape:\n",
        "        # [batch size, 1, emb size] or [1, emb_size]\n",
        "        if tf_mask is not None:\n",
        "            tf_output = self.tf(tf_input, tgt_emb, src_key_padding_mask=tf_mask)\n",
        "        else:\n",
        "            tf_output = self.tf(tf_input, tgt_emb)\n",
        "        tf_output = tf_output.reshape([-1, tf_output.shape[-1]])\n",
        "\n",
        "        # Project down to tf_to_fc_dim\n",
        "        return self.fc2(tf_output)\n",
        "\n",
        "    def _process_obs(self, obs):\n",
        "        new_obs = {}\n",
        "\n",
        "        for key in obs.keys():\n",
        "            if key not in [self.tf_input_key, self.tf_mask_key]:\n",
        "                new_obs[key] = obs[key]\n",
        "\n",
        "        tf_input = obs[self.tf_input_key]\n",
        "\n",
        "        # pass tf_input through transformer to produce embeddings\n",
        "        if self.tf_mask_key is not None:\n",
        "            tf_mask = obs[self.tf_mask_key]\n",
        "            new_obs[self.tf_input_key] = self._process_seq(tf_input, tf_mask)\n",
        "        else:\n",
        "            new_obs[self.tf_input_key] = self._process_seq(tf_input)\n",
        "\n",
        "        return new_obs\n",
        "\n",
        "    def forward(self, obs, deterministic: bool = False):\n",
        "        # Get the transformer input sequence from the obs\n",
        "        new_obs = self._process_obs(obs)\n",
        "\n",
        "        return super().forward(new_obs, deterministic)\n",
        "\n",
        "    def evaluate_actions(self, obs: torch.Tensor, actions: torch.Tensor):\n",
        "        new_obs = self._process_obs(obs)\n",
        "        return super().evaluate_actions(new_obs, actions)\n",
        "\n",
        "    def get_distribution(self, obs: torch.Tensor):\n",
        "        new_obs = self._process_obs(obs)\n",
        "        return super().get_distribution(new_obs)\n",
        "\n",
        "    def predict_values(self, obs: torch.Tensor):\n",
        "        new_obs = self._process_obs(obs)\n",
        "        return super().predict_values(new_obs)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import functools\n",
        "policy_class = functools.partial(\n",
        "        TfInputActorCriticPolicy,\n",
        "        tf_input_key=\"seq\",\n",
        "        tf_mask_key=\"seq_mask\",\n",
        "        tf_to_fc_dim=128,\n",
        "        raw_seq_dim=vlm_dim,\n",
        "        tf_kwargs={\n",
        "            \"d_model\": 512, # Set transformer token sizes to 512 dim\n",
        "            \"nhead\": 2,\n",
        "            \"num_encoder_layers\": 1,\n",
        "            \"num_decoder_layers\": 1,\n",
        "            \"dim_feedforward\": 512,\n",
        "            \"batch_first\": True,\n",
        "        },\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we use initialize the model and training algorithm, then run training with the standard procedure set out by Stable Baselines3 reinforcement learning tool suite."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from stable_baselines3 import PPO\n",
        "model = PPO(\n",
        "        policy_class,\n",
        "        env,\n",
        "        learning_rate=1e-4,\n",
        "        n_steps=64,\n",
        "        device=\"cuda:0\",\n",
        "        ent_coef=5e-3,\n",
        "        verbose=1,\n",
        "        stats_window_size=1000,\n",
        "        policy_kwargs={\n",
        "            \"net_arch\": [\n",
        "                128,\n",
        "            ]\n",
        "        },\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "model.learn(\n",
        "        total_timesteps=20480,\n",
        "        progress_bar=True,\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from stable_baselines3.common.env_util import make_vec_env\n",
        "from tqdm import tqdm\n",
        "\n",
        "# Parallel environments\n",
        "vec_env = make_vec_env(env)\n",
        "\n",
        "total_successes = 0\n",
        "for i in tqdm(range(1000)):\n",
        "    obs = vec_env.reset()\n",
        "    action, _states = model.predict(obs)\n",
        "    obs, rewards, dones, info = vec_env.step(action)\n",
        "    total_successes += rewards"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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",
      "version": "3.10.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}