action_decoder_eval.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/levi/code/flava-action-xformer/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "# Now lets connect the current model to RLBench to step through the task and evaluate how often it receives a reward in the task.\n",
    "from rlbench.action_modes.action_mode import MoveArmThenGripper\n",
    "from rlbench.action_modes.arm_action_modes import JointVelocity\n",
    "from rlbench.action_modes.gripper_action_modes import Discrete\n",
    "from rlbench.environment import Environment\n",
    "from rlbench.observation_config import ObservationConfig\n",
    "from rlbench.tasks import FS10_V1, ReachTarget\n",
    "\n",
    "from transformers import FlavaProcessor, FlavaModel\n",
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from action_decoder_model import ActionDecoderModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_PATH = \"/home/levi/data/action_decoder_model_20231024-115056.pt\"\n",
    "levi = torch.load(MODEL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['pos_encoder.pe', 'transformer_decoder.layers.0.self_attn.in_proj_weight', 'transformer_decoder.layers.0.self_attn.in_proj_bias', 'transformer_decoder.layers.0.self_attn.out_proj.weight', 'transformer_decoder.layers.0.self_attn.out_proj.bias', 'transformer_decoder.layers.0.multihead_attn.in_proj_weight', 'transformer_decoder.layers.0.multihead_attn.in_proj_bias', 'transformer_decoder.layers.0.multihead_attn.out_proj.weight', 'transformer_decoder.layers.0.multihead_attn.out_proj.bias', 'transformer_decoder.layers.0.linear1.weight', 'transformer_decoder.layers.0.linear1.bias', 'transformer_decoder.layers.0.linear2.weight', 'transformer_decoder.layers.0.linear2.bias', 'transformer_decoder.layers.0.norm1.weight', 'transformer_decoder.layers.0.norm1.bias', 'transformer_decoder.layers.0.norm2.weight', 'transformer_decoder.layers.0.norm2.bias', 'transformer_decoder.layers.0.norm3.weight', 'transformer_decoder.layers.0.norm3.bias', 'transformer_decoder.layers.1.self_attn.in_proj_weight', 'transformer_decoder.layers.1.self_attn.in_proj_bias', 'transformer_decoder.layers.1.self_attn.out_proj.weight', 'transformer_decoder.layers.1.self_attn.out_proj.bias', 'transformer_decoder.layers.1.multihead_attn.in_proj_weight', 'transformer_decoder.layers.1.multihead_attn.in_proj_bias', 'transformer_decoder.layers.1.multihead_attn.out_proj.weight', 'transformer_decoder.layers.1.multihead_attn.out_proj.bias', 'transformer_decoder.layers.1.linear1.weight', 'transformer_decoder.layers.1.linear1.bias', 'transformer_decoder.layers.1.linear2.weight', 'transformer_decoder.layers.1.linear2.bias', 'transformer_decoder.layers.1.norm1.weight', 'transformer_decoder.layers.1.norm1.bias', 'transformer_decoder.layers.1.norm2.weight', 'transformer_decoder.layers.1.norm2.bias', 'transformer_decoder.layers.1.norm3.weight', 'transformer_decoder.layers.1.norm3.bias', 'linear_action_in.weight', 'linear_action_in.bias', 'linear_action_out.weight', 'linear_action_out.bias'])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "levi.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "class Agent(object):\n",
    "\n",
    "    def __init__(self, action_shape, model_path):\n",
    "        self.action_shape = action_shape        \n",
    "        self.encoder_emb = []\n",
    "        # begin with the sos token\n",
    "        sos = np.zeros(self.action_shape, dtype=np.float32)\n",
    "        sos[0::2] = -1 # even values are -1\n",
    "        self.decoder_actions = [sos]\n",
    "        # Retrieve the Flava model and processor\n",
    "        self.flava_model = FlavaModel.from_pretrained('facebook/flava-full')\n",
    "        self.flava_processor = FlavaProcessor.from_pretrained('facebook/flava-full')\n",
    "        # Retrieve the saved decoder model\n",
    "        self.action_decoder_model = torch.load(model_path)\n",
    "\n",
    "\n",
    "    def get_flava_embeddings(self, img, instruction):\n",
    "        # Convert the observation and instruction into a batch of inputs for the Flava model\n",
    "        inputs = self.flava_processor(img, instruction, return_tensors=\"pt\", padding=\"max_length\", max_length=197, return_codebook_pixels=False, return_image_mask=False)\n",
    "        # Pass the inputs through the Flava model\n",
    "        outputs = self.flava_model(**inputs)\n",
    "        # Retrieve the multimodal embeddings from the Flava model outputs\n",
    "        multimodal_embeddings = outputs.multimodal_embeddings.detach().numpy()\n",
    "        return multimodal_embeddings\n",
    "    \n",
    "    def act(self, img, instruction):        \n",
    "        # Get the Flava embeddings for the observation and instruction\n",
    "        encoder_emb = self.get_flava_embeddings(img, instruction)\n",
    "        # Apply mean pooling to the encoder embeddings to get a single embedding for the observation\n",
    "        self.encoder_emb.append(np.mean(encoder_emb, axis=1))        \n",
    "        self.action_decoder_model.eval()  # turn on evaluation mode    \n",
    "        with torch.no_grad():\n",
    "            # Get the decoder action from the action decoder model\n",
    "            decoder_action = self.action_decoder_model(actions=self.decoder_actions, memory=self.encoder_emb)\n",
    "            # Get the action from the decoder output\n",
    "            action = decoder_action[0, -1, :].detach().numpy()\n",
    "            # Add the action to the decoder actions\n",
    "            self.decoder_actions.append(action)\n",
    "        return action\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_config = ObservationConfig()\n",
    "obs_config.set_all(True)\n",
    "\n",
    "env = Environment(\n",
    "    action_mode=MoveArmThenGripper(\n",
    "        arm_action_mode=JointVelocity(), gripper_action_mode=Discrete()),\n",
    "    obs_config=ObservationConfig(),\n",
    "    headless=True)\n",
    "env.launch()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`text_config_dict` is provided which will be used to initialize `FlavaTextConfig`. The value `text_config[\"id2label\"]` will be overriden.\n",
      "`multimodal_config_dict` is provided which will be used to initialize `FlavaMultimodalConfig`. The value `multimodal_config[\"id2label\"]` will be overriden.\n",
      "`image_codebook_config_dict` is provided which will be used to initialize `FlavaImageCodebookConfig`. The value `image_codebook_config[\"id2label\"]` will be overriden.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at facebook/flava-full were not used when initializing FlavaModel: ['mmm_image_head.transform.dense.weight', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_4.weight', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_3.weight', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_3.weight', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_1.bias', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_3.bias', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_1.bias', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_2.bias', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_3.weight', 'image_codebook.blocks.group_2.group.block_1.id_path.weight', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_3.bias', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_4.bias', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_1.weight', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_4.weight', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_1.weight', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_2.weight', 'image_codebook.blocks.input.weight', 'image_codebook.blocks.input.bias', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_3.weight', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_4.bias', 'itm_head.pooler.dense.bias', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_2.bias', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_3.bias', 'itm_head.seq_relationship.weight', 'image_codebook.blocks.output.conv.weight', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_4.weight', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_2.bias', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_4.bias', 'image_codebook.blocks.group_3.group.block_1.id_path.weight', 'mmm_image_head.decoder.bias', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_2.bias', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_2.weight', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_4.weight', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_4.bias', 'mim_head.decoder.weight', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_2.weight', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_1.weight', 'mmm_image_head.transform.LayerNorm.bias', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_2.bias', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_2.weight', 'mlm_head.transform.LayerNorm.bias', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_1.bias', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_3.weight', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_3.bias', 'image_codebook.blocks.output.conv.bias', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_1.weight', 'mmm_text_head.bias', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_3.bias', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_1.bias', 'mlm_head.decoder.bias', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_4.weight', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_1.bias', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_3.bias', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_1.weight', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_1.bias', 'mim_head.transform.dense.bias', 'mlm_head.transform.dense.weight', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_2.weight', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_4.weight', 'mmm_text_head.transform.LayerNorm.weight', 'mmm_image_head.decoder.weight', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_3.weight', 'itm_head.seq_relationship.bias', 'mmm_image_head.transform.LayerNorm.weight', 'mlm_head.transform.LayerNorm.weight', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_3.weight', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_4.weight', 'image_codebook.blocks.group_4.group.block_1.res_path.path.conv_1.weight', 'mim_head.decoder.bias', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_1.bias', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_2.bias', 'mmm_image_head.transform.dense.bias', 'image_codebook.blocks.group_2.group.block_2.res_path.path.conv_3.bias', 'itm_head.pooler.dense.weight', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_4.bias', 'mlm_head.decoder.weight', 'mmm_image_head.bias', 'mlm_head.bias', 'image_codebook.blocks.group_4.group.block_1.id_path.bias', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_3.weight', 'mmm_text_head.transform.LayerNorm.bias', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_1.weight', 'mmm_text_head.transform.dense.weight', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_4.bias', 'mim_head.bias', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_2.weight', 'mlm_head.transform.dense.bias', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_2.bias', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_4.bias', 'image_codebook.blocks.group_3.group.block_1.res_path.path.conv_2.weight', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_4.bias', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_1.weight', 'mmm_text_head.transform.dense.bias', 'image_codebook.blocks.group_2.group.block_1.id_path.bias', 'mmm_text_head.decoder.bias', 'image_codebook.blocks.group_3.group.block_2.res_path.path.conv_2.bias', 'image_codebook.blocks.group_3.group.block_1.id_path.bias', 'image_codebook.blocks.group_4.group.block_2.res_path.path.conv_1.bias', 'mim_head.transform.LayerNorm.weight', 'mim_head.transform.LayerNorm.bias', 'image_codebook.blocks.group_1.group.block_2.res_path.path.conv_3.bias', 'image_codebook.blocks.group_2.group.block_1.res_path.path.conv_4.weight', 'mim_head.transform.dense.weight', 'image_codebook.blocks.group_4.group.block_1.id_path.weight', 'mmm_text_head.decoder.weight', 'image_codebook.blocks.group_1.group.block_1.res_path.path.conv_2.weight']\n",
      "- This IS expected if you are initializing FlavaModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing FlavaModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "MODEL_PATH = \"/home/levi/data/action_decoder_model_20231024-115056.pt\"\n",
    "\n",
    "# Instantiate the agent\n",
    "agent = Agent(env.action_shape, MODEL_PATH)\n",
    "# Get the task\n",
    "task = env.get_task(ReachTarget)\n",
    "task.sample_variation()  # random variation\n",
    "# Reset the task\n",
    "descriptions, obs = task.reset()\n",
    "instruction = descriptions[1] # Could make this random at some point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/levi/code/flava-action-xformer/.venv/lib/python3.11/site-packages/transformers/modeling_utils.py:884: FutureWarning: The `device` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'collections.OrderedDict' object has no attribute 'eval'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)\n",
      "\u001b[1;32m/home/levi/code/flava-action-xformer/main.ipynb Cell 20\u001b[0m line \u001b[0;36m2\n",
      "\u001b[1;32m      <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m \u001b[39m# We can execute this cell multiple times to step through the task\u001b[39;00m\n",
      "\u001b[0;32m----> <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=1'>2</a>\u001b[0m action \u001b[39m=\u001b[39m agent\u001b[39m.\u001b[39;49mact(obs\u001b[39m.\u001b[39;49mfront_rgb, instruction)\n",
      "\n",
      "\u001b[1;32m/home/levi/code/flava-action-xformer/main.ipynb Cell 20\u001b[0m line \u001b[0;36m3\n",
      "\u001b[1;32m     <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=32'>33</a>\u001b[0m \u001b[39m# Apply mean pooling to the encoder embeddings to get a single embedding for the observation\u001b[39;00m\n",
      "\u001b[1;32m     <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=33'>34</a>\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mencoder_emb\u001b[39m.\u001b[39mappend(np\u001b[39m.\u001b[39mmean(encoder_emb, axis\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m))        \n",
      "\u001b[0;32m---> <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=34'>35</a>\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49maction_decoder_model\u001b[39m.\u001b[39;49meval()  \u001b[39m# turn on evaluation mode    \u001b[39;00m\n",
      "\u001b[1;32m     <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=35'>36</a>\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mno_grad():\n",
      "\u001b[1;32m     <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=36'>37</a>\u001b[0m     \u001b[39m# Get the decoder action from the action decoder model\u001b[39;00m\n",
      "\u001b[1;32m     <a href='vscode-notebook-cell:/home/levi/code/flava-action-xformer/main.ipynb#X30sZmlsZQ%3D%3D?line=37'>38</a>\u001b[0m     decoder_action \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39maction_decoder_model(actions\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdecoder_actions, memory\u001b[39m=\u001b[39m\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mencoder_emb)\n",
      "\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'collections.OrderedDict' object has no attribute 'eval'"
     ]
    }
   ],
   "source": [
    "# We can execute this cell multiple times to step through the task\n",
    "action = agent.act(obs.front_rgb, instruction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "obs, reward, terminate = task.step(action)\n",
    "print('Step: {} Reward: {}'.format(i, reward))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}