ViT_Mamba_QAgent.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Preprocessed state shape: torch.Size([1, 3, 224, 224]), type: <class 'torch.Tensor'>\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 0\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 1\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 2\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 3\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 4\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 5\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 6\n",
      "Vision Transformer output shape: torch.Size([1, 197, 768]), type: <class 'torch.Tensor'>\n",
      "Input to QNetwork shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Inside QNetwork, input shape: torch.Size([1, 768]), type: <class 'torch.Tensor'>\n",
      "Buffer size before sampling: 7\n",
      "Sampled states shape: torch.Size([8, 3, 224, 224]), type: <class 'torch.Tensor'>\n",
      "states types: <built-in method type of Tensor object at 0x0000028AB71A5F10>, states shape: torch.Size([8, 3, 224, 224])\n",
      "next_states types: <built-in method type of Tensor object at 0x0000028AB71A5F70>, next_states shape: torch.Size([8, 3, 224, 224])\n",
      "actions types: <built-in method type of Tensor object at 0x0000028AB71A5CD0>, actions shape: torch.Size([8, 1])\n",
      "rewards types: <built-in method type of Tensor object at 0x0000028AB71A5C10>, rewards shape: torch.Size([8])\n",
      "dones types: <built-in method type of Tensor object at 0x0000028AB71A5D30>, dones shape: torch.Size([8])\n",
      "weights types: <built-in method type of Tensor object at 0x0000028AB71A5D90>, weights shape: torch.Size([8])\n",
      "Inside QNetwork, input shape: torch.Size([8, 3, 224, 224]), type: <class 'torch.Tensor'>\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "mat1 and mat2 shapes cannot be multiplied (5376x224 and 768x128)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[21], line 563\u001b[0m\n\u001b[0;32m    561\u001b[0m     \u001b[38;5;66;03m# Update Q-network periodically\u001b[39;00m\n\u001b[0;32m    562\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(replay_buffer\u001b[38;5;241m.\u001b[39mbuffer) \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m batch_size:\n\u001b[1;32m--> 563\u001b[0m         \u001b[43mupdate_q_network\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmamba_model\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mq_network\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget_q_network\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreplay_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0.99\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Assume n_step=3 and gamma=0.99\u001b[39;00m\n\u001b[0;32m    565\u001b[0m     preprocessed_state \u001b[38;5;241m=\u001b[39m next_preprocessed_state\n\u001b[0;32m    566\u001b[0m state \u001b[38;5;241m=\u001b[39m env\u001b[38;5;241m.\u001b[39mreset()\n",
      "Cell \u001b[1;32mIn[21], line 369\u001b[0m, in \u001b[0;36mupdate_q_network\u001b[1;34m(model, q_network, target_q_network, buffer, optimizer, batch_size, gamma, n_step)\u001b[0m\n\u001b[0;32m    367\u001b[0m \u001b[38;5;66;03m# Calculate the target Q-values\u001b[39;00m\n\u001b[0;32m    368\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[1;32m--> 369\u001b[0m     next_q_values \u001b[38;5;241m=\u001b[39m \u001b[43mtarget_q_network\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnext_states\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    370\u001b[0m     max_next_q_values \u001b[38;5;241m=\u001b[39m next_q_values\u001b[38;5;241m.\u001b[39mmax(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)[\u001b[38;5;241m0\u001b[39m]\n\u001b[0;32m    371\u001b[0m     q_targets \u001b[38;5;241m=\u001b[39m rewards \u001b[38;5;241m+\u001b[39m (gamma \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m n_step) \u001b[38;5;241m*\u001b[39m max_next_q_values \u001b[38;5;241m*\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m-\u001b[39m dones)\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[21], line 279\u001b[0m, in \u001b[0;36mQNetwork.forward\u001b[1;34m(self, state_representation)\u001b[0m\n\u001b[0;32m    276\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, state_representation):\n\u001b[0;32m    277\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mInside QNetwork, input shape: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mstate_representation\u001b[38;5;241m.\u001b[39mshape\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, type: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(state_representation)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m--> 279\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinear1\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstate_representation\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    280\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrelu(x)\n\u001b[0;32m    281\u001b[0m     q_values \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlinear2(x)\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\linear.py:114\u001b[0m, in \u001b[0;36mLinear.forward\u001b[1;34m(self, input)\u001b[0m\n\u001b[0;32m    113\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[1;32m--> 114\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinear\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbias\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mRuntimeError\u001b[0m: mat1 and mat2 shapes cannot be multiplied (5376x224 and 768x128)"
     ]
    }
   ],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import copy\n",
    "\n",
    "import gym\n",
    "from gym import spaces\n",
    "# Assuming you have PIL (Python Imaging Library) installed for image processing\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "#VIT\n",
    "class FlexiblePatchEmbedding(nn.Module):\n",
    "    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_size=768, temporal_patch_size=1, is_3d=False):\n",
    "        super().__init__()\n",
    "        self.img_size = img_size\n",
    "        self.patch_size = patch_size\n",
    "        self.in_channels = in_channels\n",
    "        self.embed_size = embed_size\n",
    "        self.is_3d = is_3d\n",
    "\n",
    "        if is_3d:\n",
    "            self.num_patches = int(((img_size // patch_size) ** 2) * temporal_patch_size)\n",
    "            self.projection = nn.Conv3d(in_channels, embed_size, kernel_size=(temporal_patch_size, patch_size, patch_size), \n",
    "\n",
    "                                        stride=(temporal_patch_size, patch_size, patch_size))\n",
    "\n",
    "        else:\n",
    "            self.num_patches = (img_size // patch_size) ** 2\n",
    "            self.projection = nn.Conv2d(in_channels, embed_size, kernel_size=patch_size, stride=patch_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.projection(x)  # [B, E, T/P, H/P, W/P] or [B, E, H/P, W/P]\n",
    "        x = x.flatten(2)  # Flatten spatial and temporal dimensions\n",
    "        x = x.transpose(1, 2)  # [B, N, E]\n",
    "        return x\n",
    "\n",
    "class PositionalEmbedding(nn.Module):\n",
    "    def __init__(self, num_patches, embed_size):\n",
    "        super().__init__()\n",
    "\n",
    "        self.positional_embedding = nn.Parameter(torch.zeros(1, num_patches + 1, embed_size))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        cls_token = torch.zeros(batch_size, 1, x.shape[-1], device=x.device)\n",
    "        x = torch.cat([cls_token, x], dim=1)  # [B, 1+N, E]\n",
    "        x += self.positional_embedding\n",
    "        return x\n",
    "\n",
    "class TransformerEncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):\n",
    "        super(TransformerEncoderLayer, self).__init__()\n",
    "        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)\n",
    "\n",
    "        # Feedforward network\n",
    "        self.linear1 = nn.Linear(d_model, dim_feedforward)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear2 = nn.Linear(dim_feedforward, d_model)\n",
    "\n",
    "        self.norm1 = nn.LayerNorm(d_model)\n",
    "        self.norm2 = nn.LayerNorm(d_model)\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "\n",
    "        self.activation = nn.ReLU()\n",
    "\n",
    "    def forward(self, src):\n",
    "        src2 = self.self_attn(src, src, src)[0]\n",
    "        src = src + self.dropout1(src2)\n",
    "        src = self.norm1(src)\n",
    "\n",
    "        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))\n",
    "        src = src + self.dropout2(src2)\n",
    "        src = self.norm2(src)\n",
    "        return src\n",
    "\n",
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, encoder_layer, num_layers):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)])\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "    def forward(self, src):\n",
    "        for layer in self.layers:\n",
    "            src = layer(src)\n",
    "\n",
    "        return src\n",
    "\n",
    "\n",
    "class FlexibleVisionTransformer(nn.Module):\n",
    "    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_size=768, num_heads=12, num_layers=12, num_classes=1000, temporal_patch_size=1, is_3d=False):\n",
    "        super().__init__()\n",
    "        self.patch_embedding = FlexiblePatchEmbedding(img_size, patch_size, in_channels, embed_size, temporal_patch_size, is_3d)\n",
    "        self.positional_embedding = PositionalEmbedding(self.patch_embedding.num_patches, embed_size)\n",
    "\n",
    "        encoder_layer = TransformerEncoderLayer(d_model=embed_size, nhead=num_heads)\n",
    "        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_layers)\n",
    "\n",
    "        self.mlp_head = nn.Sequential(\n",
    "            nn.LayerNorm(embed_size),\n",
    "            nn.Linear(embed_size, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, return_embeddings=False):\n",
    "        x = self.patch_embedding(x)\n",
    "        x = self.positional_embedding(x)\n",
    "        x = self.transformer_encoder(x)\n",
    "        if return_embeddings:\n",
    "            return x  # Return the sequence of embeddings directly\n",
    "        cls_token = x[:, 0]\n",
    "        x = self.mlp_head(cls_token)\n",
    "        return x\n",
    "\n",
    "\n",
    "# DualStateMamba\n",
    "class MambaModelWithDualState(nn.Module):\n",
    "    def __init__(self, d_model, merge_interval=10):\n",
    "        super().__init__()\n",
    "        self.merge_interval = merge_interval\n",
    "        self.d_model = d_model\n",
    "\n",
    "        # Simulated components for updating and merging states\n",
    "        self.state_updater = nn.Linear(d_model, d_model)\n",
    "        self.state_merger = nn.Linear(2 * d_model, d_model)  # Merges two states into one\n",
    "        self.state_projector = nn.Linear(d_model, d_model)  # Projects merged state to next state\n",
    "\n",
    "        # Initial states\n",
    "        self.reset_states()\n",
    "\n",
    "    def reset_states(self):\n",
    "        # Resets/initializes the states; assumes states are kept between forward passes\n",
    "        self.vit_state = torch.zeros(1, self.d_model)\n",
    "        self.q_network_state = torch.zeros(1, self.d_model)\n",
    "\n",
    "    def update_state(self, current_state, new_input):\n",
    "        if new_input is None:\n",
    "            # If new_input is None, return the current state without change\n",
    "            return current_state\n",
    "        else:\n",
    "            # Simple state update mechanism; in practice, this could involve more complex temporal processing\n",
    "            return F.relu(self.state_updater(new_input + current_state))\n",
    "\n",
    "    def merge_states(self, vit_state, q_network_state):\n",
    "        # Merges two states; this could be a complex learned fusion operation\n",
    "        merged = torch.cat([vit_state, q_network_state], dim=1)\n",
    "        return F.relu(self.state_merger(merged))\n",
    "\n",
    "    def project_to_next_state(self, merged_state):\n",
    "        # Projects the merged state to the next state; could involve temporal dynamics modeling\n",
    "        return F.relu(self.state_projector(merged_state))\n",
    "\n",
    "    def forward(self, vit_features, q_network_action, timestep):\n",
    "        # Update states based on new inputs\n",
    "        self.vit_state = self.update_state(self.vit_state, vit_features)\n",
    "        self.q_network_state = self.update_state(self.q_network_state, q_network_action)\n",
    "\n",
    "        if timestep % self.merge_interval == 0:\n",
    "            # Periodic merge operation\n",
    "            merged_state = self.merge_states(self.vit_state, self.q_network_state)\n",
    "            next_state = self.project_to_next_state(merged_state)\n",
    "\n",
    "            # Update both states with the merged and projected state\n",
    "            self.vit_state, self.q_network_state = next_state.clone(), next_state.clone()\n",
    "\n",
    "        # Optionally, return current state for integration with other components\n",
    "\n",
    "        return self.vit_state, self.q_network_state\n",
    "\n",
    "\n",
    "\n",
    "# Example usage\n",
    "d_model = 512  # Feature dimension\n",
    "mamba_model = MambaModelWithDualState(d_model=d_model)\n",
    "\n",
    "# Simulated inputs for demonstration\n",
    "vit_features = torch.randn(1, d_model)  # Output from ViT for current frame\n",
    "q_network_action = torch.randn(1, d_model)  # Representation of action selected by Q-Network\n",
    "timestep = 1  # Current timestep\n",
    "\n",
    "# Forward pass through the model\n",
    "vit_state, q_network_state = mamba_model(vit_features, q_network_action, timestep)\n",
    "\n",
    "# The returned states (vit_state, q_network_state) could be used for further processing, decision making, etc.\n",
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, d_model, eps=1e-8):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.eps = eps\n",
    "        self.scale = nn.Parameter(torch.ones(d_model))\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean_square = torch.mean(x**2, dim=-1, keepdim=True)\n",
    "        rms = torch.sqrt(mean_square + self.eps)\n",
    "        return x / rms * self.scale\n",
    "\n",
    "class MambaBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, expand_factor=2, d_state=512):\n",
    "        super().__init__()\n",
    "        self.d_inner = expand_factor * d_model\n",
    "        self.conv1d = nn.Conv1d(in_channels=d_model, out_channels=self.d_inner, kernel_size=3, padding=1)\n",
    "        self.x_proj = nn.Linear(self.d_inner, d_state)\n",
    "        self.out_proj = nn.Linear(d_state, d_model)\n",
    "        self.silu = nn.SiLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.transpose(1, 2)  # Assuming [batch, feature, seq_len] for Conv1D\n",
    "        x, z = torch.chunk(self.conv1d(x), 2, dim=1)\n",
    "        x = x.transpose(1, 2)  # Switch back to [batch, seq_len, feature]\n",
    "        ssm_params = self.x_proj(self.silu(x))\n",
    "        z = self.silu(z.transpose(1, 2))\n",
    "\n",
    "        # Assuming a selective_scan or equivalent function exists and operates on ssm_params\n",
    "        # This is a placeholder to represent some form of state-space computation\n",
    "        x = ssm_params * z  # Simplified representation of a state-space operation\n",
    "        out = self.out_proj(x)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, d_model):\n",
    "        super().__init__()\n",
    "        self.norm = RMSNorm(d_model)\n",
    "        self.mamba_block = MambaBlock(d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.mamba_block(self.norm(x))\n",
    "        return x + out\n",
    "\n",
    "\n",
    "\n",
    "class MambaModel(nn.Module):\n",
    "    def __init__(self, d_model, n_layers):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([ResidualBlock(d_model) for _ in range(n_layers)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        return x\n",
    "\n",
    "def selective_scan(delta, B, x):\n",
    "    \"\"\"Implements the selective scan operation.\n",
    "    Args:\n",
    "        delta (Tensor): Delta vector (shape: [batch_size, d_state])\n",
    "        B (Tensor): B matrix (shape: [batch_size, d_state, d_state])\n",
    "        x (Tensor): Input (shape: [batch_size, seq_len, d_state + d_input])\n",
    "\n",
    "    Returns:\n",
    "        Tensor: Updated state representation (shape: [batch_size, seq_len, d_state])\n",
    "    \"\"\"\n",
    "    batch_size, seq_len, _ = x.shape\n",
    "    state = torch.zeros(batch_size, d_state).to(x.device) \n",
    "    output = torch.zeros(batch_size, seq_len, d_state).to(x.device)\n",
    "\n",
    "    for t in range(seq_len):\n",
    "        xt = x[:, t, :]  # Current input with state concatenated \n",
    "        state = B[:, t] * state  + xt * delta[:, t]  \n",
    "        output[:, t, :] = state\n",
    "\n",
    "    return output\n",
    "\n",
    "\n",
    "# Q Network\n",
    "class QNetwork(nn.Module):\n",
    "    def __init__(self, mamba_output_dim, num_actions):\n",
    "        super().__init__()\n",
    "        self.linear1 = nn.Linear(mamba_output_dim, 128)  \n",
    "        self.relu = nn.ReLU()\n",
    "        self.linear2 = nn.Linear(128, num_actions)  \n",
    "\n",
    "    def forward(self, state_representation):\n",
    "        print(f'Inside QNetwork, input shape: {state_representation.shape}, type: {type(state_representation)}')\n",
    "\n",
    "        x = self.linear1(state_representation)\n",
    "        x = self.relu(x)\n",
    "        q_values = self.linear2(x)\n",
    "        return q_values\n",
    "\n",
    "\n",
    "\n",
    "class ReplayBuffer:\n",
    "    def __init__(self, capacity, n_step, gamma):\n",
    "        self.capacity = capacity\n",
    "        self.buffer = []  # List to store experiences\n",
    "        self.priorities = SumTree(capacity)  # Use a SumTree for efficient priority sampling\n",
    "        self.n_step = n_step\n",
    "        self.gamma = gamma \n",
    "\n",
    "    def _calculate_n_step_return(self, start_idx):\n",
    "        \"\"\"Calculates the n-step return for an experience.\"\"\"\n",
    "        reward = 0\n",
    "\n",
    "        for i in range(self.n_step):\n",
    "            idx = start_idx + i\n",
    "            if idx >= len(self.buffer):\n",
    "                break\n",
    "            reward += self.buffer[idx][2] * self.gamma**i  # Discounted reward\n",
    "            if self.buffer[idx][4]:  # Done flag\n",
    "                break\n",
    "        return reward \n",
    "\n",
    "    def add(self, state, action, reward, next_state, done):\n",
    "        \"\"\"Adds a new experience and updates priorities.\"\"\"\n",
    "        experience = (state, action, reward, next_state, done)\n",
    "        max_priority = self.priorities.max() or 1.0  # Avoid 0 priority\n",
    "        self.buffer.append(experience)\n",
    "        self.priorities.add(max_priority, len(self.buffer) - 1)\n",
    "\n",
    "\n",
    "    def sample(self, batch_size, beta=0.4):\n",
    "        \"\"\"Samples a batch of experiences with importance-sampling weights.\"\"\"\n",
    "        indices = self.priorities.sample(batch_size)\n",
    "        weights = np.zeros(batch_size, dtype=np.float32)\n",
    "        batch = [self.buffer[idx] for idx in indices]\n",
    "        \n",
    "        total_priority = self.priorities.total_priority()\n",
    "        min_priority = self.priorities.min() / total_priority\n",
    "        max_weight = (min_priority * batch_size) ** (-beta)\n",
    "        \n",
    "        for i, idx in enumerate(indices):\n",
    "            priority = self.priorities.get(idx)[0] / total_priority\n",
    "            weight = (priority * batch_size) ** (-beta) / max_weight\n",
    "            weights[i] = weight\n",
    "        \n",
    "        states, actions, rewards, next_states, dones = zip(*batch)\n",
    "        \n",
    "        # Convert lists to tensors. Note: Assuming your states and next_states are images and already in the correct shape.\n",
    "        states = torch.stack([torch.Tensor(state) for state in states])\n",
    "        next_states = torch.stack([torch.Tensor(state) for state in next_states])\n",
    "        actions = torch.tensor(actions, dtype=torch.long).unsqueeze(1)\n",
    "        rewards = torch.tensor(rewards, dtype=torch.float32)\n",
    "        dones = torch.tensor(dones, dtype=torch.float32)\n",
    "        print(f'Sampled states shape: {states.shape}, type: {type(states)}')\n",
    "\n",
    "        return states, actions, rewards, next_states, dones, indices, torch.tensor(weights, dtype=torch.float32)\n",
    "\n",
    "\n",
    "    def update_priorities(self, indices, td_errors):\n",
    "        \"\"\"Updates priorities based on TD errors.\"\"\"\n",
    "        for idx, td_error in zip(indices, td_errors):\n",
    "            priority = abs(td_error) + 1e-5  # Avoid zero priority\n",
    "            self.priorities.update(idx, priority)\n",
    "\n",
    "\n",
    "\n",
    "def update_q_network(model, q_network, target_q_network, buffer, optimizer, batch_size, gamma, n_step):\n",
    "    # Sample a batch of experiences\n",
    "    states, actions, rewards, next_states, dones, indices, weights = buffer.sample(batch_size)\n",
    "    print(f\"states types: {states.type}, states shape: {states.shape}\")\n",
    "    print(f\"next_states types: {next_states.type}, next_states shape: {next_states.shape}\")\n",
    "    print(f\"actions types: {actions.type}, actions shape: {actions.shape}\")\n",
    "    print(f\"rewards types: {rewards.type}, rewards shape: {rewards.shape}\")\n",
    "    print(f\"dones types: {dones.type}, dones shape: {dones.shape}\")\n",
    "    #print(f\"indices types: {indices.type}\")\n",
    "    print(f\"weights types: {weights.type}, weights shape: {weights.shape}\")\n",
    "\n",
    "    actions = actions.unsqueeze(1) if actions.dim() == 1 else actions\n",
    "    rewards = rewards if rewards.dim() > 0 else torch.tensor(rewards, dtype=torch.float32)\n",
    "    dones = dones if dones.dim() > 0 else torch.tensor(dones, dtype=torch.float32)\n",
    "    weights = weights if weights.dim() > 0 else torch.tensor(weights, dtype=torch.float32)\n",
    "\n",
    "    # Calculate the target Q-values\n",
    "    with torch.no_grad():\n",
    "        next_q_values = target_q_network(next_states)\n",
    "        max_next_q_values = next_q_values.max(dim=1)[0]\n",
    "        q_targets = rewards + (gamma ** n_step) * max_next_q_values * (1 - dones)\n",
    "\n",
    "    # Calculate the current Q-values\n",
    "    current_q_values = q_network(states).gather(1, actions).squeeze()\n",
    "\n",
    "    # Compute loss using the weights from the prioritized replay buffer\n",
    "    loss = (weights * F.mse_loss(current_q_values, q_targets, reduction='none')).mean()\n",
    "\n",
    "    # Perform an optimization step\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    # Update the priorities in the buffer\n",
    "    td_errors = torch.abs(current_q_values - q_targets).detach().numpy()\n",
    "    buffer.update_priorities(indices, td_errors)\n",
    "\n",
    "\n",
    "\n",
    "class SumTree:\n",
    "    def __init__(self, capacity):\n",
    "        self.capacity = capacity\n",
    "        self.tree = np.zeros(2 * capacity - 1)  # Binary tree, twice the capacity for leaves and parents\n",
    "        self.data_pointer = 0  # Pointer to the next free data index\n",
    "\n",
    "    def add(self, priority, data_index):\n",
    "        tree_index = self.data_pointer + self.capacity - 1\n",
    "        self.update(tree_index, priority)  # Update the new leaf's value\n",
    "        self.data_pointer += 1\n",
    "\n",
    "        if self.data_pointer >= self.capacity: \n",
    "            self.data_pointer = 0  # Wrap around if the buffer is full\n",
    "\n",
    "    def update(self, tree_index, priority):\n",
    "        change = priority - self.tree[tree_index]\n",
    "        self.tree[tree_index] = priority\n",
    "        self._propagate_changes(tree_index, change)  # Update values up the tree\n",
    "\n",
    "    def min(self):\n",
    "        \"\"\"Return the smallest non-zero priority.\"\"\"\n",
    "        min_val = np.min(self.tree[-self.capacity:][self.tree[-self.capacity:] > 0])\n",
    "        if min_val == 0:\n",
    "            return np.min(self.tree[-self.capacity:])  # Fallback to min, including zeros\n",
    "        return min_val\n",
    "\n",
    "    def _propagate_changes(self, tree_index, change):\n",
    "        parent_index = (tree_index - 1) // 2\n",
    "        self.tree[parent_index] += change\n",
    "\n",
    "        if parent_index != 0:\n",
    "            self._propagate_changes(parent_index, change)\n",
    "\n",
    "    def get(self, sample):\n",
    "        parent_index = 0\n",
    "        while True:\n",
    "            left_child_index = 2 * parent_index + 1\n",
    "            right_child_index = left_child_index + 1\n",
    "\n",
    "            if left_child_index >= len(self.tree):  # Leaf node\n",
    "                leaf_index = parent_index\n",
    "                break\n",
    "\n",
    "            else:  # Continue searching\n",
    "                if sample <= self.tree[left_child_index]:\n",
    "                    parent_index = left_child_index\n",
    "                else:\n",
    "                    sample -= self.tree[left_child_index]\n",
    "                    parent_index = right_child_index\n",
    "\n",
    "        data_index = leaf_index - self.capacity + 1\n",
    "        return leaf_index, self.tree[leaf_index], data_index\n",
    "\n",
    "    def total_priority(self):\n",
    "        return self.tree[0]  # Root node holds the sum of all priorities\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        batch_indices = []\n",
    "        p_segments = self.total_priority() / batch_size\n",
    "\n",
    "        for i in range(batch_size):\n",
    "            start = i * p_segments\n",
    "            end = (i + 1) * p_segments\n",
    "            sample = random.uniform(start, end)\n",
    "            _, _, idx = self.get(sample)\n",
    "            batch_indices.append(idx)\n",
    "\n",
    "        return batch_indices\n",
    "\n",
    "    def max(self):\n",
    "        return np.max(self.tree[-self.capacity:])  # Max priority among leave\n",
    "    \n",
    "\n",
    "\n",
    "\n",
    "class VisualEnvironment(gym.Env):\n",
    "    def __init__(self):\n",
    "        super(VisualEnvironment, self).__init__()\n",
    "        self.action_space = spaces.Discrete(2)  # Example: left or right\n",
    "        self.observation_space = spaces.Box(low=0, high=255, shape=(224, 224, 3), dtype=np.uint8)\n",
    "        self.state = None\n",
    "\n",
    "    def step(self, action):\n",
    "        reward = 0\n",
    "        done = False\n",
    "        info = {}\n",
    "        # Generate a random image for the next state\n",
    "        self.state = np.random.randint(0, 256, (224, 224, 3), dtype=np.uint8)\n",
    "        return self.state, reward, done, info\n",
    "\n",
    "    def reset(self):\n",
    "        # Generate an initial random image\n",
    "        self.state = np.random.randint(0, 256, (224, 224, 3), dtype=np.uint8)\n",
    "        return self.state\n",
    "\n",
    "\n",
    "\n",
    "def preprocess_state(image, target_size=(224, 224)):\n",
    "    \"\"\"\n",
    "    Preprocess the state image for the Vision Transformer model.\n",
    "    Args:\n",
    "        image (np.ndarray): The state image as a NumPy array.\n",
    "        target_size (tuple): The target image size (width, height).\n",
    "    Returns:\n",
    "        torch.Tensor: The preprocessed image tensor.\n",
    "    \"\"\"\n",
    "    transform = transforms.Compose([\n",
    "        transforms.ToPILImage(),  # Convert np.ndarray to PIL.Image\n",
    "        transforms.Resize(target_size),  # Resize to target size\n",
    "        transforms.ToTensor(),  # Convert PIL.Image to torch.Tensor\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406],  # Standard normalization for ImageNet\n",
    "                             std=[0.229, 0.224, 0.225])\n",
    "    ])\n",
    "    image_tensor = transform(image)\n",
    "    # Add a batch dimension (B, C, H, W) where B=1 since you're processing single images\n",
    "    image_tensor = image_tensor.unsqueeze(0)\n",
    "    return image_tensor\n",
    "\n",
    "# Assuming an environment class Env that supports env.reset() and env.step(action)\n",
    "env = VisualEnvironment()\n",
    "state = env.reset()  # Get initial state as a raw image\n",
    "preprocessed_state = preprocess_state(state)  # Preprocess the initial state\n",
    "done = False\n",
    "\n",
    "# Instantiate the model, Q-network, and other components\n",
    "vit_model = FlexibleVisionTransformer(\n",
    "    img_size=224,\n",
    "    patch_size=16,\n",
    "    in_channels=3,\n",
    "    embed_size=768,\n",
    "    num_heads=12,\n",
    "    num_layers=12,\n",
    "    num_classes=1000,\n",
    "    temporal_patch_size=1,\n",
    "    is_3d=False)\n",
    "\n",
    "vit_features = vit_model(preprocessed_state)\n",
    "mamba_model = MambaModelWithDualState(d_model=768)\n",
    "q_network = QNetwork(mamba_output_dim=768, num_actions=env.action_space.n)\n",
    "target_q_network = copy.deepcopy(q_network)  # For stable Q-value targets\n",
    "replay_buffer = ReplayBuffer(capacity=10000, n_step=3, gamma=0.99)\n",
    "optimizer = torch.optim.Adam(q_network.parameters(), lr=1e-4)\n",
    "total_episodes = 100\n",
    "batch_size = 8\n",
    "# Example training loop\n",
    "for episode in range(total_episodes):\n",
    "    state = env.reset()  # Get initial state as a raw image\n",
    "    preprocessed_state = preprocess_state(state)  # Preprocess the initial state\n",
    "    print(f'Preprocessed state shape: {preprocessed_state.shape}, type: {type(preprocessed_state)}')\n",
    "\n",
    "    done = False\n",
    "    \n",
    "    while not done:\n",
    "        # vit_model and mamba_model accept the preprocessed state directly\n",
    "        vit_output = vit_model(preprocessed_state, return_embeddings=True)\n",
    "        print(f'Vision Transformer output shape: {vit_output.shape}, type: {type(vit_output)}')\n",
    "\n",
    "        cls_token = vit_output[:, 0, :]  # Extract the cls_token\n",
    "        vit_state, q_network_state = mamba_model(cls_token, None, timestep)        \n",
    "        print(f'Input to QNetwork shape: {vit_state.shape}, type: {type(vit_state)}')\n",
    "\n",
    "\n",
    "        action = q_network(vit_state).argmax().item()  # Decide action\n",
    "        next_raw_state, reward, done, _ = env.step(action)  # Execute action in env\n",
    "        next_preprocessed_state = preprocess_state(next_raw_state)  # Preprocess the new state\n",
    "\n",
    "        print(f'Buffer size before sampling: {len(replay_buffer.buffer)}')\n",
    "\n",
    "        # Store the transition in replay buffer\n",
    "        replay_buffer.add(preprocessed_state.squeeze(0).cpu().numpy(), action, reward, next_preprocessed_state.squeeze(0).cpu().numpy(), done)\n",
    "        \n",
    "        # Update Q-network periodically\n",
    "        if len(replay_buffer.buffer) >= batch_size:\n",
    "            update_q_network(mamba_model, q_network, target_q_network, replay_buffer, optimizer, batch_size, 0.99, 3)  # Assume n_step=3 and gamma=0.99\n",
    "        \n",
    "        preprocessed_state = next_preprocessed_state\n",
    "    state = env.reset()\n",
    "    done = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training\n",
    "1. **Initialization:**\n",
    "   - Instantiate the environment, models (VIT, MambaModel, QNetwork, and Target QNetwork), replay buffer, and optimizer.\n",
    "   - Set hyperparameters (learning rate, batch size, discount factor gamma, update frequency for the target network).\n",
    "\n",
    "2. **Episode Loop:**\n",
    "   - Reset the environment to get the initial state.\n",
    "\n",
    "3. **Timestep Loop within an Episode:**\n",
    "   - Process the current state through VIT and optionally MambaModel to get a state representation.\n",
    "   - Select an action based on the current state representation.\n",
    "   - Execute the action in the environment to obtain the next state, reward, and done flag.\n",
    "   - Preprocess the next state and add the transition to the replay buffer.\n",
    "   - Sample a batch from the replay buffer and perform a gradient descent step on the Q-network.\n",
    "   - Every fixed number of steps, update the target Q-network's weights to match the current Q-network's weights.\n",
    "   - Check for the episode termination condition. If true, break the loop.\n",
    "\n",
    "4. **Replay Buffer:**\n",
    "   - Ensure transitions stored and sampled include preprocessed states or appropriate state representations, not raw pixel data.\n",
    "\n",
    "5. **Q-Network Optimization:**\n",
    "   - Calculate expected Q values from the Q-network using sampled states.\n",
    "   - Calculate target Q values using the next state's max Q value from the target Q-network.\n",
    "   - Compute the loss between expected and target Q values.\n",
    "   - Perform a backpropagation and optimizer step to update the Q-network's weights.\n",
    "\n",
    "6. **Target Network Update:**\n",
    "   - Periodically sync the target Q-network's weights with the current Q-network's weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(224, 224, 3)\n",
      "Batch size: 3\n",
      "Input shape to FlexiblePatchEmbedding: torch.Size([3, 224, 224])\n",
      "Shape after projection: torch.Size([768, 14, 14])\n",
      "Number of patches received for positional embeddings: 14\n",
      "Shape of x before CLS token: torch.Size([768, 14, 14])\n",
      "x shape: torch.Size([768, 15, 14])\n",
      "Positional embedding shape: torch.Size([1, 769, 768])\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "The size of tensor a (14) must match the size of tensor b (768) at non-singleton dimension 2",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[73], line 668\u001b[0m\n\u001b[0;32m    655\u001b[0m \u001b[38;5;66;03m# Instantiate the model, Q-network, and other components\u001b[39;00m\n\u001b[0;32m    656\u001b[0m \u001b[38;5;66;03m# Instantiate the model, Q-network, and other components\u001b[39;00m\n\u001b[0;32m    657\u001b[0m vit_model \u001b[38;5;241m=\u001b[39m FlexibleVisionTransformer(\n\u001b[0;32m    658\u001b[0m     img_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m224\u001b[39m,\n\u001b[0;32m    659\u001b[0m     patch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m16\u001b[39m,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    664\u001b[0m     temporal_patch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[0;32m    665\u001b[0m     is_3d\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[1;32m--> 668\u001b[0m vit_features \u001b[38;5;241m=\u001b[39m \u001b[43mvit_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpreprocessed_state\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    669\u001b[0m mamba_model \u001b[38;5;241m=\u001b[39m MambaModelWithDualState(d_model\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m768\u001b[39m)\n\u001b[0;32m    670\u001b[0m q_network \u001b[38;5;241m=\u001b[39m QNetwork(mamba_output_dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m768\u001b[39m, num_actions\u001b[38;5;241m=\u001b[39menv\u001b[38;5;241m.\u001b[39maction_space\u001b[38;5;241m.\u001b[39mn)\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[73], line 132\u001b[0m, in \u001b[0;36mFlexibleVisionTransformer.forward\u001b[1;34m(self, x, return_embeddings)\u001b[0m\n\u001b[0;32m    129\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBatch size: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mx\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m) \n\u001b[0;32m    131\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpatch_embedding(x)\n\u001b[1;32m--> 132\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpositional_embedding\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    133\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtransformer_encoder(x)\n\u001b[0;32m    134\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m return_embeddings:\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1517\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1518\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\robbi\\anaconda3\\envs\\my_gpu_env_llm\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1522\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1523\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1525\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1526\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1527\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1529\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1530\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "Cell \u001b[1;32mIn[73], line 70\u001b[0m, in \u001b[0;36mPositionalEmbedding.forward\u001b[1;34m(self, x)\u001b[0m\n\u001b[0;32m     65\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPositional embedding shape: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpositional_embedding\u001b[38;5;241m.\u001b[39mshape\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     67\u001b[0m \u001b[38;5;66;03m# Ensure that the positional_embedding is correctly sized to match `x`\u001b[39;00m\n\u001b[0;32m     68\u001b[0m \u001b[38;5;66;03m#assert x.shape[1] == self.positional_embedding.shape[1], \"Mismatch in positional embedding size\"\u001b[39;00m\n\u001b[1;32m---> 70\u001b[0m \u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpositional_embedding\u001b[49m\n\u001b[0;32m     71\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m x\n",
      "\u001b[1;31mRuntimeError\u001b[0m: The size of tensor a (14) must match the size of tensor b (768) at non-singleton dimension 2"
     ]
    }
   ],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import copy\n",
    "\n",
    "import gym\n",
    "from gym import spaces\n",
    "# Assuming you have PIL (Python Imaging Library) installed for image processing\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "from collections import namedtuple\n",
    "\n",
    "\n",
    "\n",
    "#VIT\n",
    "class FlexiblePatchEmbedding(nn.Module):\n",
    "    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_size=768, temporal_patch_size=1, is_3d=False):\n",
    "        super().__init__()\n",
    "        self.img_size = img_size\n",
    "        self.patch_size = patch_size\n",
    "        self.in_channels = in_channels\n",
    "        self.embed_size = embed_size\n",
    "        self.is_3d = is_3d\n",
    "\n",
    "        if is_3d:\n",
    "            self.num_patches = int(((img_size // patch_size) ** 2) * temporal_patch_size)\n",
    "            self.projection = nn.Conv3d(in_channels, embed_size, kernel_size=(temporal_patch_size, patch_size, patch_size), \n",
    "\n",
    "                                        stride=(temporal_patch_size, patch_size, patch_size))\n",
    "\n",
    "        else:\n",
    "            (img_size // patch_size) ** 2 * embed_size            \n",
    "\n",
    "            self.projection = nn.Conv2d(in_channels, embed_size, kernel_size=patch_size, stride=patch_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        print(f\"Input shape to FlexiblePatchEmbedding: {x.shape}\")\n",
    "\n",
    "        x = self.projection(x)  # [B, E, T/P, H/P, W/P] or [B, E, H/P, W/P]\n",
    "        print(f\"Shape after projection: {x.shape}\")\n",
    "\n",
    "        x = x.transpose(1, 2)  # [B, N, E]  – Keep spatial information \n",
    "        x = x.flatten(2)\n",
    "        return x \n",
    "\n",
    "class PositionalEmbedding(nn.Module):\n",
    "    def __init__(self, num_patches, embed_size):\n",
    "        super().__init__()\n",
    "\n",
    "        # Initialize positional embeddings with an extra position for the cls_token\n",
    "        # *** Key Change: Create the positional embedding directly ***\n",
    "        self.positional_embedding = nn.Parameter(torch.zeros(1, embed_size + 1, embed_size)) \n",
    "\n",
    "    def forward(self, x):\n",
    "        print(f\"Number of patches received for positional embeddings: {x.shape[1]}\")\n",
    "        print(f\"Shape of x before CLS token: {x.shape}\") \n",
    "        batch_size = x.shape[0]\n",
    "        cls_token = torch.zeros(batch_size, 1, x.shape[-1], device=x.device)\n",
    "        x = torch.cat([cls_token, x], dim=1)  # [B, 1+N, E]\n",
    "\n",
    "        # Debugging: Print shapes to ensure compatibility\n",
    "        print(f\"x shape: {x.shape}\")\n",
    "        print(f\"Positional embedding shape: {self.positional_embedding.shape}\")\n",
    "\n",
    "        # Ensure that the positional_embedding is correctly sized to match `x`\n",
    "        #assert x.shape[1] == self.positional_embedding.shape[1], \"Mismatch in positional embedding size\"\n",
    "        \n",
    "        x += self.positional_embedding\n",
    "        return x\n",
    "\n",
    "\n",
    "class TransformerEncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):\n",
    "        super(TransformerEncoderLayer, self).__init__()\n",
    "        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)\n",
    "\n",
    "        # Feedforward network\n",
    "        self.linear1 = nn.Linear(d_model, dim_feedforward)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear2 = nn.Linear(dim_feedforward, d_model)\n",
    "\n",
    "        self.norm1 = nn.LayerNorm(d_model)\n",
    "        self.norm2 = nn.LayerNorm(d_model)\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "\n",
    "        self.activation = nn.ReLU()\n",
    "\n",
    "    def forward(self, src):\n",
    "        src2 = self.self_attn(src, src, src)[0]\n",
    "        src = src + self.dropout1(src2)\n",
    "        src = self.norm1(src)\n",
    "\n",
    "        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))\n",
    "        src = src + self.dropout2(src2)\n",
    "        src = self.norm2(src)\n",
    "        return src\n",
    "\n",
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, encoder_layer, num_layers):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)])\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "    def forward(self, src):\n",
    "        for layer in self.layers:\n",
    "            src = layer(src)\n",
    "\n",
    "        return src\n",
    "\n",
    "\n",
    "class FlexibleVisionTransformer(nn.Module):\n",
    "    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_size=768, num_heads=12, num_layers=12, num_classes=1000, temporal_patch_size=1, is_3d=False):\n",
    "        super().__init__()\n",
    "        self.patch_embedding = FlexiblePatchEmbedding(img_size, patch_size, in_channels, embed_size, temporal_patch_size, is_3d)\n",
    "        self.positional_embedding = PositionalEmbedding(embed_size, embed_size) \n",
    "\n",
    "        encoder_layer = TransformerEncoderLayer(d_model=embed_size, nhead=num_heads)\n",
    "        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_layers)\n",
    "\n",
    "        self.mlp_head = nn.Sequential(\n",
    "            nn.LayerNorm(embed_size),\n",
    "            nn.Linear(embed_size, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x, return_embeddings=False):\n",
    "        print(f\"Batch size: {x.shape[0]}\") \n",
    "\n",
    "        x = self.patch_embedding(x)\n",
    "        x = self.positional_embedding(x)\n",
    "        x = self.transformer_encoder(x)\n",
    "        if return_embeddings:\n",
    "            return x  # Return the sequence of embeddings directly\n",
    "        cls_token = x[:, 0]\n",
    "        x = self.mlp_head(cls_token)\n",
    "        return x\n",
    "\n",
    "\n",
    "# DualStateMamba\n",
    "class MambaModelWithDualState(nn.Module):\n",
    "    def __init__(self, d_model, merge_interval=10):\n",
    "        super().__init__()\n",
    "        self.merge_interval = merge_interval\n",
    "        self.d_model = d_model\n",
    "\n",
    "        # Simulated components for updating and merging states\n",
    "        self.state_updater = nn.Linear(d_model, d_model)\n",
    "        self.state_merger = nn.Linear(2 * d_model, d_model)  # Merges two states into one\n",
    "        self.state_projector = nn.Linear(d_model, d_model)  # Projects merged state to next state\n",
    "\n",
    "        # Initial states\n",
    "        self.reset_states()\n",
    "\n",
    "    def reset_states(self):\n",
    "        # Resets/initializes the states; assumes states are kept between forward passes\n",
    "        self.vit_state = torch.zeros(1, self.d_model)\n",
    "        self.q_network_state = torch.zeros(1, self.d_model)\n",
    "\n",
    "    def update_state(self, current_state, new_input):\n",
    "        if new_input is None:\n",
    "            # If new_input is None, return the current state without change\n",
    "            return current_state\n",
    "        else:\n",
    "            # Simple state update mechanism; in practice, this could involve more complex temporal processing\n",
    "            return F.relu(self.state_updater(new_input + current_state))\n",
    "\n",
    "    def merge_states(self, vit_state, q_network_state):\n",
    "        # Merges two states; this could be a complex learned fusion operation\n",
    "        merged = torch.cat([vit_state, q_network_state], dim=1)\n",
    "        return F.relu(self.state_merger(merged))\n",
    "\n",
    "    def project_to_next_state(self, merged_state):\n",
    "        # Projects the merged state to the next state; could involve temporal dynamics modeling\n",
    "        return F.relu(self.state_projector(merged_state))\n",
    "\n",
    "    def forward(self, vit_features, q_network_action, timestep):\n",
    "        # Update states based on new inputs\n",
    "        self.vit_state = self.update_state(self.vit_state, vit_features)\n",
    "        self.q_network_state = self.update_state(self.q_network_state, q_network_action)\n",
    "\n",
    "        if timestep % self.merge_interval == 0:\n",
    "            # Periodic merge operation\n",
    "            merged_state = self.merge_states(self.vit_state, self.q_network_state)\n",
    "            next_state = self.project_to_next_state(merged_state)\n",
    "\n",
    "            # Update both states with the merged and projected state\n",
    "            self.vit_state, self.q_network_state = next_state.clone(), next_state.clone()\n",
    "\n",
    "        # Optionally, return current state for integration with other components\n",
    "\n",
    "        return self.vit_state, self.q_network_state\n",
    "\n",
    "# The returned states (vit_state, q_network_state) could be used for further processing, decision making, etc.\n",
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, d_model, eps=1e-8):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.eps = eps\n",
    "        self.scale = nn.Parameter(torch.ones(d_model))\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean_square = torch.mean(x**2, dim=-1, keepdim=True)\n",
    "        rms = torch.sqrt(mean_square + self.eps)\n",
    "        return x / rms * self.scale\n",
    "\n",
    "class MambaBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, expand_factor=2, d_state=512):\n",
    "        super().__init__()\n",
    "        self.d_inner = expand_factor * d_model\n",
    "        self.conv1d = nn.Conv1d(in_channels=d_model, out_channels=self.d_inner, kernel_size=3, padding=1)\n",
    "        self.x_proj = nn.Linear(self.d_inner, d_state)\n",
    "        self.out_proj = nn.Linear(d_state, d_model)\n",
    "        self.silu = nn.SiLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.transpose(1, 2)  # Assuming [batch, feature, seq_len] for Conv1D\n",
    "        x, z = torch.chunk(self.conv1d(x), 2, dim=1)\n",
    "        x = x.transpose(1, 2)  # Switch back to [batch, seq_len, feature]\n",
    "        ssm_params = self.x_proj(self.silu(x))\n",
    "        z = self.silu(z.transpose(1, 2))\n",
    "\n",
    "        # Assuming a selective_scan or equivalent function exists and operates on ssm_params\n",
    "        # This is a placeholder to represent some form of state-space computation\n",
    "        x = ssm_params * z  # Simplified representation of a state-space operation\n",
    "        out = self.out_proj(x)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, d_model):\n",
    "        super().__init__()\n",
    "        self.norm = RMSNorm(d_model)\n",
    "        self.mamba_block = MambaBlock(d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.mamba_block(self.norm(x))\n",
    "        return x + out\n",
    "\n",
    "class MambaModel(nn.Module):\n",
    "    def __init__(self, d_model, n_layers):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([ResidualBlock(d_model) for _ in range(n_layers)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        return x\n",
    "\n",
    "def selective_scan(delta, B, x):\n",
    "    \"\"\"Implements the selective scan operation.\n",
    "    Args:\n",
    "        delta (Tensor): Delta vector (shape: [batch_size, d_state])\n",
    "        B (Tensor): B matrix (shape: [batch_size, d_state, d_state])\n",
    "        x (Tensor): Input (shape: [batch_size, seq_len, d_state + d_input])\n",
    "\n",
    "    Returns:\n",
    "        Tensor: Updated state representation (shape: [batch_size, seq_len, d_state])\n",
    "    \"\"\"\n",
    "    batch_size, seq_len, _ = x.shape\n",
    "    state = torch.zeros(batch_size, d_state).to(x.device) \n",
    "    output = torch.zeros(batch_size, seq_len, d_state).to(x.device)\n",
    "\n",
    "    for t in range(seq_len):\n",
    "        xt = x[:, t, :]  # Current input with state concatenated \n",
    "        state = B[:, t] * state  + xt * delta[:, t]  \n",
    "        output[:, t, :] = state\n",
    "\n",
    "    return output\n",
    "\n",
    "\n",
    "# Env\n",
    "\n",
    "\n",
    "class VisualEnvironment(gym.Env):\n",
    "    def __init__(self):\n",
    "        super(VisualEnvironment, self).__init__()\n",
    "        self.action_space = spaces.Discrete(2)  # Example: left or right\n",
    "        self.observation_space = spaces.Box(low=0, high=255, shape=(224, 224, 3), dtype=np.uint8)\n",
    "        self.state = None\n",
    "\n",
    "    def step(self, action):\n",
    "        reward = 0\n",
    "        done = False\n",
    "        info = {}\n",
    "        # Generate a random image for the next state\n",
    "        self.state = np.random.randint(0, 256, (224, 224, 3), dtype=np.uint8)\n",
    "        return self.state, reward, done, info\n",
    "\n",
    "    def reset(self):\n",
    "        # Generate an initial random image\n",
    "        self.state = np.random.randint(0, 256, (224, 224, 3), dtype=np.uint8)\n",
    "        return self.state\n",
    "\n",
    "\n",
    "\n",
    "def preprocess_state(image, target_size=(224, 224)):\n",
    "    \"\"\"\n",
    "    Preprocess the state image for the Vision Transformer model.\n",
    "    Args:\n",
    "        image (np.ndarray or torch.Tensor): The state image as a NumPy array or a PyTorch tensor.\n",
    "        target_size (tuple): The target image size (width, height).\n",
    "    Returns:\n",
    "        torch.Tensor: The preprocessed image tensor.\n",
    "    \"\"\"\n",
    "    if isinstance(image, np.ndarray):\n",
    "        # Convert NumPy array to torch tensor\n",
    "        image_tensor = torch.tensor(image)\n",
    "        # Check if the image has more than 3 dimensions (e.g., batch dimension)\n",
    "        if len(image_tensor.shape) > 3:\n",
    "            # Remove the batch dimension if present\n",
    "            image_tensor = image_tensor.squeeze(0)\n",
    "        # Permute dimensions to match PyTorch convention (assuming the input is HWC)\n",
    "        image_tensor = image_tensor.permute(2, 0, 1)\n",
    "    elif isinstance(image, torch.Tensor):\n",
    "        # If already a torch tensor, no need for conversion\n",
    "        image_tensor = image\n",
    "    else:\n",
    "        raise TypeError(\"Input type should be np.ndarray or torch.Tensor\")\n",
    "\n",
    "    # Resize and normalize the image tensor\n",
    "    transform = transforms.Compose([\n",
    "        transforms.ToPILImage(),  # Convert to PIL Image\n",
    "        transforms.Resize(target_size),  # Resize to target size\n",
    "        transforms.ToTensor(),  # Convert to tensor\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) \n",
    "    ])\n",
    "\n",
    "    image_tensor = transform(image_tensor)  # Apply transformations\n",
    "\n",
    "    # Add a batch dimension if it's not present\n",
    "    if len(image_tensor.shape) < 3:\n",
    "        image_tensor = image_tensor.unsqueeze(0)\n",
    "\n",
    "    return image_tensor\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Q Network\n",
    "class QNetwork(nn.Module):\n",
    "    def __init__(self, mamba_output_dim, num_actions):\n",
    "        super().__init__()\n",
    "        self.linear1 = nn.Linear(mamba_output_dim, 128)  \n",
    "        self.relu = nn.ReLU()\n",
    "        self.linear2 = nn.Linear(128, num_actions)  \n",
    "\n",
    "    def forward(self, state_representation):\n",
    "        print(f'Inside QNetwork, input shape: {state_representation.shape}, type: {type(state_representation)}')\n",
    "\n",
    "        x = self.linear1(state_representation)\n",
    "        x = self.relu(x)\n",
    "        q_values = self.linear2(x)\n",
    "        return q_values\n",
    "\n",
    "\n",
    "class SumTree:\n",
    "    def __init__(self, capacity):\n",
    "        self.capacity = capacity\n",
    "        self.tree = np.zeros(2 * capacity - 1)  # Binary tree, twice the capacity for leaves and parents\n",
    "        self.data_pointer = 0  # Pointer to the next free data index\n",
    "\n",
    "    def add(self, priority, data_index):\n",
    "        tree_index = self.data_pointer + self.capacity - 1\n",
    "        self.update(tree_index, priority)  # Update the new leaf's value\n",
    "        self.data_pointer += 1\n",
    "\n",
    "        if self.data_pointer >= self.capacity: \n",
    "            self.data_pointer = 0  # Wrap around if the buffer is full\n",
    "\n",
    "    def update(self, tree_index, priority):\n",
    "        change = priority - self.tree[tree_index]\n",
    "        self.tree[tree_index] = priority\n",
    "        self._propagate_changes(tree_index, change)  # Update values up the tree\n",
    "\n",
    "    def min(self):\n",
    "        \"\"\"Return the smallest non-zero priority.\"\"\"\n",
    "        min_val = np.min(self.tree[-self.capacity:][self.tree[-self.capacity:] > 0])\n",
    "        if min_val == 0:\n",
    "            return np.min(self.tree[-self.capacity:])  # Fallback to min, including zeros\n",
    "        return min_val\n",
    "\n",
    "    def _propagate_changes(self, tree_index, change):\n",
    "        parent_index = (tree_index - 1) // 2\n",
    "        self.tree[parent_index] += change\n",
    "\n",
    "        if parent_index != 0:\n",
    "            self._propagate_changes(parent_index, change)\n",
    "\n",
    "    def get(self, sample):\n",
    "        parent_index = 0\n",
    "        while True:\n",
    "            left_child_index = 2 * parent_index + 1\n",
    "            right_child_index = left_child_index + 1\n",
    "\n",
    "            if left_child_index >= len(self.tree):  # Leaf node\n",
    "                leaf_index = parent_index\n",
    "                break\n",
    "\n",
    "            else:  # Continue searching\n",
    "                if sample <= self.tree[left_child_index]:\n",
    "                    parent_index = left_child_index\n",
    "                else:\n",
    "                    sample -= self.tree[left_child_index]\n",
    "                    parent_index = right_child_index\n",
    "\n",
    "        data_index = leaf_index - self.capacity + 1\n",
    "        return leaf_index, self.tree[leaf_index], data_index\n",
    "\n",
    "    def total_priority(self):\n",
    "        return self.tree[0]  # Root node holds the sum of all priorities\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        batch_indices = []\n",
    "        p_segments = self.total_priority() / batch_size\n",
    "\n",
    "        for i in range(batch_size):\n",
    "            start = i * p_segments\n",
    "            end = (i + 1) * p_segments\n",
    "            sample = random.uniform(start, end)\n",
    "            _, _, idx = self.get(sample)\n",
    "            batch_indices.append(idx)\n",
    "\n",
    "        return batch_indices\n",
    "\n",
    "    def max(self):\n",
    "        return np.max(self.tree[-self.capacity:])  # Max priority among leave\n",
    "    \n",
    "\n",
    "class ReplayBuffer:\n",
    "    def __init__(self, capacity, n_step, gamma):\n",
    "        self.capacity = capacity\n",
    "        self.buffer = []  # List to store experiences\n",
    "        self.priorities = SumTree(capacity)  # Use a SumTree for efficient priority sampling\n",
    "        self.n_step = n_step\n",
    "        self.gamma = gamma \n",
    "\n",
    "    def _calculate_n_step_return(self, start_idx):\n",
    "        \"\"\"Calculates the n-step return for an experience.\"\"\"\n",
    "        reward = 0\n",
    "\n",
    "        for i in range(self.n_step):\n",
    "            idx = start_idx + i\n",
    "            if idx >= len(self.buffer):\n",
    "                break\n",
    "            reward += self.buffer[idx][2] * self.gamma**i  # Discounted reward\n",
    "            if self.buffer[idx][4]:  # Done flag\n",
    "                break\n",
    "        return reward \n",
    "\n",
    "    def add(self, state, action, reward, next_state, done):\n",
    "        \"\"\"Adds a new experience and updates priorities.\"\"\"\n",
    "        experience = (state, action, reward, next_state, done)\n",
    "        max_priority = self.priorities.max() or 1.0  # Avoid 0 priority\n",
    "        self.buffer.append(experience)\n",
    "        self.priorities.add(max_priority, len(self.buffer) - 1)\n",
    "\n",
    "\n",
    "    def sample(self, batch_size, beta=0.4):\n",
    "        \"\"\"Samples a batch of experiences with importance-sampling weights.\"\"\"\n",
    "        indices = self.priorities.sample(batch_size)\n",
    "        weights = np.zeros(batch_size, dtype=np.float32)\n",
    "        batch = [self.buffer[idx] for idx in indices]\n",
    "        \n",
    "        total_priority = self.priorities.total_priority()\n",
    "        min_priority = self.priorities.min() / total_priority\n",
    "        max_weight = (min_priority * batch_size) ** (-beta)\n",
    "        \n",
    "        for i, idx in enumerate(indices):\n",
    "            priority = self.priorities.get(idx)[0] / total_priority\n",
    "            weight = (priority * batch_size) ** (-beta) / max_weight\n",
    "            weights[i] = weight\n",
    "        \n",
    "        states, actions, rewards, next_states, dones = zip(*batch)\n",
    "        \n",
    "        # Convert lists to tensors. Note: Assuming your states and next_states are images and already in the correct shape.\n",
    "        states = torch.stack([torch.Tensor(state) for state in states])\n",
    "        next_states = torch.stack([torch.Tensor(state) for state in next_states])\n",
    "        actions = torch.tensor(actions, dtype=torch.long).unsqueeze(1)\n",
    "        rewards = torch.tensor(rewards, dtype=torch.float32)\n",
    "        dones = torch.tensor(dones, dtype=torch.float32)\n",
    "        print(f'Sampled states shape: {states.shape}, type: {type(states)}')\n",
    "\n",
    "        return states, actions, rewards, next_states, dones, indices, torch.tensor(weights, dtype=torch.float32)\n",
    "\n",
    "\n",
    "    def update_priorities(self, indices, td_errors):\n",
    "        \"\"\"Updates priorities based on TD errors.\"\"\"\n",
    "        for idx, td_error in zip(indices, td_errors):\n",
    "            priority = abs(td_error) + 1e-5  # Avoid zero priority\n",
    "            self.priorities.update(idx, priority)\n",
    "\n",
    "\n",
    "def update_q_network(q_network, target_q_network, replay_buffer, optimizer, gamma, batch_size, n_step):\n",
    "    # Sample a batch of experiences if the buffer is large enough\n",
    "    if len(replay_buffer.buffer) < batch_size:\n",
    "        return\n",
    "\n",
    "    states, actions, rewards, next_states, dones, indices, weights = replay_buffer.sample(batch_size)\n",
    "\n",
    "    # Explicitly convert dones to a Boolean tensor\n",
    "    dones = dones.bool()\n",
    "\n",
    "    non_final_mask = ~dones\n",
    "    non_final_next_states = next_states[non_final_mask]\n",
    "\n",
    "    state_action_values = q_network(states).gather(1, actions.view(-1, 1)).squeeze()\n",
    "\n",
    "    next_state_values = torch.zeros(batch_size)\n",
    "    if non_final_next_states.size(0) > 0:  # Check if there are any non-final next states\n",
    "        next_state_values[non_final_mask] = target_q_network(non_final_next_states).max(1)[0].detach()\n",
    "    expected_state_action_values = (next_state_values * gamma ** n_step) + rewards\n",
    "\n",
    "    loss = (weights * F.mse_loss(state_action_values, expected_state_action_values.detach(), reduction='none')).mean()\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    td_errors = torch.abs(state_action_values - expected_state_action_values).detach().numpy()\n",
    "    replay_buffer.update_priorities(indices, td_errors)\n",
    "\n",
    "\n",
    "Transition = namedtuple('Transition', ('state', 'action', 'reward', 'next_state', 'done'))\n",
    "\n",
    "class TrainingManager:\n",
    "    def __init__(self, env, vit_model, mamba_model, q_network, target_q_network, replay_buffer, optimizer, batch_size=32, gamma=0.99, update_target_every=100):\n",
    "        self.env = env\n",
    "        self.vit_model = vit_model\n",
    "        self.mamba_model = mamba_model\n",
    "        self.q_network = q_network\n",
    "        self.target_q_network = target_q_network\n",
    "        self.replay_buffer = replay_buffer\n",
    "        self.optimizer = optimizer\n",
    "        self.batch_size = batch_size\n",
    "        self.gamma = gamma\n",
    "        self.update_target_every = update_target_every\n",
    "        self.num_actions = env.action_space.n\n",
    "        self.steps_done = 0\n",
    "\n",
    "    def preprocess(self, image):\n",
    "        # Assuming preprocess_state function is defined outside this class\n",
    "        return preprocess_state(image)\n",
    "\n",
    "    def select_action(self, state):\n",
    "        sample = random.random()\n",
    "        eps_threshold = 0.05  # Fixed strategy for simplicity\n",
    "        if sample > eps_threshold:\n",
    "            with torch.no_grad():\n",
    "                # Forward pass through the vit_model to get state features\n",
    "                state_features = self.vit_model(self.preprocess(state), return_embeddings=True)\n",
    "                # Using the cls token as state representation\n",
    "                state_representation = state_features[:, 0, :]\n",
    "                # Forward pass through the q_network to get action values\n",
    "                return self.q_network(state_representation).max(1)[1].view(1, 1)\n",
    "        else:\n",
    "            return torch.tensor([[random.randrange(self.num_actions)]], dtype=torch.long)\n",
    "        \n",
    "    @staticmethod\n",
    "    def extract_features(vit_model, states):\n",
    "        \"\"\"\n",
    "        Extract features from states using the Vision Transformer model.\n",
    "        Args:\n",
    "            vit_model (nn.Module): The Vision Transformer model.\n",
    "            states (List[torch.Tensor]): List of processed state tensors.\n",
    "        Returns:\n",
    "            torch.Tensor: Stacked tensor of extracted features.\n",
    "        \"\"\"\n",
    "        features = []\n",
    "        for state in states:\n",
    "            processed_state = preprocess_state(state)  # Process state\n",
    "            with torch.no_grad():\n",
    "                feature = vit_model(processed_state).squeeze(0)  # Forward pass through Vision Transformer\n",
    "            features.append(feature)\n",
    "        return torch.stack(features)\n",
    "\n",
    "\n",
    "\n",
    "    def optimize_model(self):\n",
    "        if len(self.replay_buffer.buffer) < self.batch_size:\n",
    "            return\n",
    "        states, actions, rewards, next_states, dones, indices, weights = replay_buffer.sample(batch_size)\n",
    "        \n",
    "        # Convert states tensor to list of images\n",
    "        states_list = [state.numpy() for state in states]\n",
    "        next_states_list = [state.numpy() for state in next_states]\n",
    "        # Process states and next_states through the Vision Transformer to get feature representations\n",
    "        state_features = TrainingManager.extract_features(vit_model, states_list)\n",
    "        next_state_features = TrainingManager.extract_features(vit_model, next_states_list)\n",
    "\n",
    "        # Now pass state_features and next_state_features to the QNetwork\n",
    "        update_q_network(q_network, \n",
    "                        target_q_network, \n",
    "                        state_features, \n",
    "                        actions, \n",
    "                        rewards, \n",
    "                        next_state_features, \n",
    "                        dones, \n",
    "                        indices, \n",
    "                        weights, \n",
    "                        optimizer, \n",
    "                        self.gamma, \n",
    "                        batch_size, \n",
    "                        self.n_step)\n",
    "\n",
    "\n",
    "\n",
    "    def train(self, num_episodes):\n",
    "        for i_episode in range(num_episodes):\n",
    "            # Initialize the environment and state\n",
    "            print(f\"Episode {i_episode}\")\n",
    "            state = self.env.reset()\n",
    "            for t in range(1000):  # Maximum steps per episode\n",
    "                print(f\"Step {t} , Episode {i_episode}\")\n",
    "\n",
    "                action = self.select_action(state)\n",
    "                next_state, reward, done, _ = self.env.step(action.item())\n",
    "                print(f\"action shape: {action.shape}\")\n",
    "                print(f\"next_state shape: {next_state.shape}\")\n",
    "                print(f\"reward shape: {reward.shape}\")\n",
    "                print(f\"done shape: {done.shape}\")\n",
    "                print(f\"state shape: {state.shape}\")\n",
    "\n",
    "\n",
    "\n",
    "                if not done:\n",
    "                    next_state_processed = self.preprocess(next_state)\n",
    "                else:\n",
    "                    next_state_processed = None\n",
    "\n",
    "                # Store the transition in memory\n",
    "                self.replay_buffer.add(state, action, reward, next_state_processed, done)\n",
    "\n",
    "                # Perform one step of the optimization\n",
    "                self.optimize_model()\n",
    "                if done:\n",
    "                    break\n",
    "\n",
    "                state = next_state\n",
    "                self.steps_done += 1\n",
    "\n",
    "\n",
    "# Assuming an environment class Env that supports env.reset() and env.step(action)\n",
    "env = VisualEnvironment()\n",
    "state = env.reset()  # Get initial state as a raw image\n",
    "print(state.shape)\n",
    "preprocessed_state = preprocess_state(state)  # Preprocess the initial state\n",
    "\n",
    "done = False\n",
    "\n",
    "# Instantiate the model, Q-network, and other components\n",
    "# Instantiate the model, Q-network, and other components\n",
    "vit_model = FlexibleVisionTransformer(\n",
    "    img_size=224,\n",
    "    patch_size=16,\n",
    "    embed_size=768,\n",
    "    num_heads=12,\n",
    "    num_layers=12,\n",
    "    num_classes=1000,\n",
    "    temporal_patch_size=1,\n",
    "    is_3d=False)\n",
    "\n",
    "\n",
    "vit_features = vit_model(preprocessed_state)\n",
    "mamba_model = MambaModelWithDualState(d_model=768)\n",
    "q_network = QNetwork(mamba_output_dim=768, num_actions=env.action_space.n)\n",
    "target_q_network = copy.deepcopy(q_network)  # For stable Q-value targets\n",
    "replay_buffer = ReplayBuffer(capacity=10000, n_step=3, gamma=0.99)\n",
    "optimizer = torch.optim.Adam(q_network.parameters(), lr=1e-4)\n",
    "total_episodes = 100\n",
    "batch_size = 8\n",
    "# Instantiate TrainingManager\n",
    "# Assuming the environment, models, replay buffer, and optimizer have been instantiated\n",
    "training_manager = TrainingManager(env, vit_model, mamba_model, q_network, target_q_network, replay_buffer, optimizer, batch_size=8)\n",
    "\n",
    "# Start training\n",
    "training_manager.train(num_episodes=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "keyword argument repeated: embed_size (2537261159.py, line 10)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;36m  Cell \u001b[1;32mIn[77], line 10\u001b[1;36m\u001b[0m\n\u001b[1;33m    pos_embed = PositionalEmbedding(embed_size=768,embed_size=768)\u001b[0m\n\u001b[1;37m                                                   ^\u001b[0m\n\u001b[1;31mSyntaxError\u001b[0m\u001b[1;31m:\u001b[0m keyword argument repeated: embed_size\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "\n",
    "# Let's simulate an image input\n",
    "image_tensor = torch.randn(3, 224, 224)  \n",
    "\n",
    "patch_embed = FlexiblePatchEmbedding(img_size=224, patch_size=16) \n",
    "pos_embed = PositionalEmbedding(embed_size=768,embed_size=768) \n",
    "\n",
    "output = patch_embed(image_tensor)\n",
    "output += pos_embed(output)  # This should work! \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "my_gpu_env_llm",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}