train_decision_transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": true,
    "id": "FJnwpdtXTuF7",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Training Decision Transformers with 🤗 transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TmRpOylHmoSo"
   },
   "source": [
    "In this tutorial, **you’ll learn to train your first Offline Decision Transformer model from scratch to make a half-cheetah run.** 🏃\n",
    "\n",
    "❓ If you have questions, please post them on #study-group discord channel 👉 https://discord.gg/aYka4Yhff9\n",
    "\n",
    "🎮 Environments:\n",
    "- [Half Cheetah](https://www.gymlibrary.dev/environments/mujoco/half_cheetah/)\n",
    "\n",
    "⬇️ Here's what you'll achieve at the end of this tutorial. ⬇️"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h78SBsjCviBm"
   },
   "source": [
    "### Prerequisites 🏗️\n",
    "Before diving into the notebook, you need to:\n",
    "\n",
    "🔲 📚 [Read the tutorial](https://huggingface.co/blog/train-decision-transformers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-21T20:20:22.981259969Z",
     "start_time": "2024-02-21T20:20:20.433204139Z"
    },
    "id": "DktITQNXTopc",
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "from dataclasses import dataclass\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from datasets import load_dataset, concatenate_datasets\n",
    "from tqdm import tqdm\n",
    "from transformers import DecisionTransformerConfig, Trainer, TrainingArguments, TrainerCallback\n",
    "from models.decision_mamba import TrainableDT, TrainableDM\n",
    "from evaluation.evaluate_episodes import evaluate_episode_rtg\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A1Ugq2POUmRA"
   },
   "source": [
    "### Step 3: Loading the dataset from the 🤗 Hub and instantiating the model\n",
    "\n",
    "We host a number of Offline RL Datasets on the hub. Today we will be training with the halfcheetah “expert” dataset, hosted here on hub.\n",
    "\n",
    "First we need to import the load_dataset function from the 🤗 datasets package and download the dataset to our machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-05T03:58:06.019369587Z",
     "start_time": "2024-01-05T03:58:03.626879617Z"
    },
    "id": "s3bLeIHqUwq7"
   },
   "outputs": [],
   "source": [
    "os.environ[\"WANDB_DISABLED\"] = \"true\" # we diable weights and biases logging for this tutorial\n",
    "DATASET_NAME = \"halfcheetah-expert-v2\"\n",
    "TOP_K = 10\n",
    "\n",
    "#dataset_medium = load_dataset(\"edbeeching/decision_transformer_gym_replay\", \"halfcheetah-medium-v2\", )\n",
    "dataset_expert = load_dataset(\"edbeeching/decision_transformer_gym_replay\", \"halfcheetah-expert-v2\", )\n",
    "\n",
    "#dataset = concatenate_datasets([dataset_medium['train'], dataset_expert['train']])\n",
    "\n",
    "dataset = dataset_expert\n",
    "\n",
    "# only save the top k trajectories and train on them\n",
    "# rewards = np.sum(dataset['train']['rewards'], axis=-1)\n",
    "# top_trajectories = np.argpartition(rewards, -TOP_K)[-TOP_K:]\n",
    "\n",
    "# dataset['train'] = dataset['train'].select(top_trajectories)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11252.035"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The highest of the dataset may not actually be what the 'max reward' from D4RL \n",
    "np.max(np.sum(dataset.with_format('np')['train']['rewards'], axis=-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZFmTdHoHUD13"
   },
   "source": [
    "### Step 4: Defining a custom DataCollator for the transformers Trainer class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-12-22T17:36:01.090703162Z",
     "start_time": "2023-12-22T17:36:01.087992422Z"
    },
    "id": "l1QzZHmPUM4p"
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class DecisionTransformerGymDataCollator:\n",
    "    return_tensors: str = \"pt\"\n",
    "    max_len: int = 20 #subsets of the episode we use for training\n",
    "    state_dim: int = 17  # size of state space\n",
    "    act_dim: int = 6  # size of action space\n",
    "    max_ep_len: int = 1000 # max episode length in the dataset\n",
    "    scale: float = 1000.0  # normalization of rewards/returns\n",
    "    state_mean: np.array = None  # to store state means\n",
    "    state_std: np.array = None  # to store state stds\n",
    "    p_sample: np.array = None  # a distribution to take account trajectory lengths\n",
    "    n_traj: int = 0 # to store the number of trajectories in the dataset\n",
    "\n",
    "    def __init__(self, dataset) -> None:\n",
    "        self.act_dim = len(dataset[0][\"actions\"][0])\n",
    "        self.state_dim = len(dataset[0][\"observations\"][0])\n",
    "        self.dataset = dataset\n",
    "        # calculate dataset stats for normalization of states\n",
    "        states = []\n",
    "        traj_lens = []\n",
    "        for obs in dataset[\"observations\"]:\n",
    "            states.extend(obs)\n",
    "            traj_lens.append(len(obs))\n",
    "        self.n_traj = len(traj_lens)\n",
    "        states = np.vstack(states)\n",
    "        self.state_mean, self.state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6\n",
    "\n",
    "        traj_lens = np.array(traj_lens)\n",
    "        self.p_sample = traj_lens / sum(traj_lens)\n",
    "\n",
    "    def _discount_cumsum(self, x, gamma):\n",
    "        discount_cumsum = np.zeros_like(x)\n",
    "        discount_cumsum[-1] = x[-1]\n",
    "        for t in reversed(range(x.shape[0] - 1)):\n",
    "            discount_cumsum[t] = x[t] + gamma * discount_cumsum[t + 1]\n",
    "        return discount_cumsum\n",
    "\n",
    "    def __call__(self, features):\n",
    "        batch_size = len(features)\n",
    "        # this is a bit of a hack to be able to sample of a non-uniform distribution\n",
    "        batch_inds = np.random.choice(\n",
    "            np.arange(self.n_traj),\n",
    "            size=batch_size,\n",
    "            replace=True,\n",
    "            p=self.p_sample,  # reweights so we sample according to timesteps\n",
    "        )\n",
    "        # a batch of dataset features\n",
    "        s, a, r, d, rtg, timesteps, mask = [], [], [], [], [], [], []\n",
    "\n",
    "        for ind in batch_inds:\n",
    "            # for feature in features:\n",
    "            feature = self.dataset[int(ind)]\n",
    "            si = random.randint(0, len(feature[\"rewards\"]) - 1)\n",
    "\n",
    "            # get sequences from dataset\n",
    "            s.append(np.array(feature[\"observations\"][si : si + self.max_len]).reshape(1, -1, self.state_dim))\n",
    "            a.append(np.array(feature[\"actions\"][si : si + self.max_len]).reshape(1, -1, self.act_dim))\n",
    "            r.append(np.array(feature[\"rewards\"][si : si + self.max_len]).reshape(1, -1, 1))\n",
    "\n",
    "            d.append(np.array(feature[\"dones\"][si : si + self.max_len]).reshape(1, -1))\n",
    "            timesteps.append(np.arange(si, si + s[-1].shape[1]).reshape(1, -1))\n",
    "            timesteps[-1][timesteps[-1] >= self.max_ep_len] = self.max_ep_len - 1  # padding cutoff\n",
    "            rtg.append(\n",
    "                self._discount_cumsum(np.array(feature[\"rewards\"][si:]), gamma=1.0)[\n",
    "                    : s[-1].shape[1]   # TODO check the +1 removed here\n",
    "                ].reshape(1, -1, 1)\n",
    "            )\n",
    "            if rtg[-1].shape[1] < s[-1].shape[1]:\n",
    "                print(\"if true\")\n",
    "                rtg[-1] = np.concatenate([rtg[-1], np.zeros((1, 1, 1))], axis=1)\n",
    "\n",
    "            # padding and state + reward normalization\n",
    "            tlen = s[-1].shape[1]\n",
    "            s[-1] = np.concatenate([np.zeros((1, self.max_len - tlen, self.state_dim)), s[-1]], axis=1)\n",
    "            s[-1] = (s[-1] - self.state_mean) / self.state_std\n",
    "            a[-1] = np.concatenate(\n",
    "                [np.ones((1, self.max_len - tlen, self.act_dim)) * -10.0, a[-1]],\n",
    "                axis=1,\n",
    "            )\n",
    "            r[-1] = np.concatenate([np.zeros((1, self.max_len - tlen, 1)), r[-1]], axis=1)\n",
    "            d[-1] = np.concatenate([np.ones((1, self.max_len - tlen)) * 2, d[-1]], axis=1)\n",
    "            rtg[-1] = np.concatenate([np.zeros((1, self.max_len - tlen, 1)), rtg[-1]], axis=1) / self.scale\n",
    "            timesteps[-1] = np.concatenate([np.zeros((1, self.max_len - tlen)), timesteps[-1]], axis=1)\n",
    "            mask.append(np.concatenate([np.zeros((1, self.max_len - tlen)), np.ones((1, tlen))], axis=1))\n",
    "\n",
    "        s = torch.from_numpy(np.concatenate(s, axis=0)).float()\n",
    "        a = torch.from_numpy(np.concatenate(a, axis=0)).float()\n",
    "        r = torch.from_numpy(np.concatenate(r, axis=0)).float()\n",
    "        d = torch.from_numpy(np.concatenate(d, axis=0))\n",
    "        rtg = torch.from_numpy(np.concatenate(rtg, axis=0)).float()\n",
    "        timesteps = torch.from_numpy(np.concatenate(timesteps, axis=0)).long()\n",
    "        mask = torch.from_numpy(np.concatenate(mask, axis=0)).float()\n",
    "\n",
    "        return {\n",
    "            \"states\": s,\n",
    "            \"actions\": a,\n",
    "            \"rewards\": r,\n",
    "            \"returns_to_go\": rtg,\n",
    "            \"timesteps\": timesteps,\n",
    "            \"attention_mask\": mask,\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bmTRGPKYUVFG"
   },
   "source": [
    "### Step 5: Extending the Decision Transformer Model to include a loss function\n",
    "\n",
    "In order to train the model with the 🤗 trainer class, we first need to ensure the dictionary it returns contains a loss, in this case L-2 norm of the models action predictions and the targets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "zIJCY3b3pQAh"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1313806"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "collator = DecisionTransformerGymDataCollator(dataset['train'])\n",
    "\n",
    "config = DecisionTransformerConfig(state_dim=collator.state_dim, act_dim=collator.act_dim)\n",
    "model = TrainableDT(config)\n",
    "\n",
    "def num_params(model):\n",
    "    return sum(p.numel() for p in model.parameters())\n",
    "\n",
    "num_params(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nJJ2mr_cU4eE"
   },
   "source": [
    "### Step 6: Defining the training hyperparameters and training the model\n",
    "Here, we define the training hyperparameters and our Trainer class that we'll use to train our Decision Transformer model.\n",
    "\n",
    "This step takes about an hour, so you may leave it running. Note the authors train for at least 3 hours, so the results presented here are not as performant as the models hosted on the 🤗 hub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "TFPuiNy-qWnP"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-0.04489212  0.03232612  0.06034821 -0.17081618 -0.19477023 -0.05751681\n",
      "  0.0970142   0.03239178 11.0473385  -0.07997213 -0.32363245  0.3629689\n",
      "  0.42323524  0.40836537  1.1085011  -0.48743752 -0.07375081]\n"
     ]
    }
   ],
   "source": [
    "import mujoco\n",
    "import gymnasium as gym\n",
    "\n",
    "# build the environment\n",
    "directory = './video'\n",
    "device = \"cuda\"\n",
    "\n",
    "model = model.to(device)\n",
    "env = gym.make(\"HalfCheetah-v4\", render_mode='rgb_array')\n",
    "\n",
    "#env = gym.wrappers.RecordVideo(env, directory)\n",
    "max_ep_len = 1000\n",
    "scale = 1000.0  # normalization for rewards/returns\n",
    "TARGET_RETURN = 12000 / scale  # evaluation is conditioned on a return of 12000, scaled accordingly\n",
    "\n",
    "state_mean = collator.state_mean.astype(np.float32)\n",
    "state_std = collator.state_std.astype(np.float32)\n",
    "print(state_mean)\n",
    "\n",
    "state_dim = env.observation_space.shape[0]\n",
    "act_dim = env.action_space.shape[0]\n",
    "# Create the decision transformer model\n",
    "\n",
    "# state_mean = torch.from_numpy(state_mean).to(device=device)\n",
    "# state_std = torch.from_numpy(state_std).to(device=device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_episodes(num_eval_episodes, model):\n",
    "    returns, lengths = [], []\n",
    "\n",
    "    model.eval()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for _ in tqdm(range(num_eval_episodes)):\n",
    "            ret, length = evaluate_episode_rtg(\n",
    "                env=env,\n",
    "                state_dim=state_dim,\n",
    "                act_dim=act_dim,\n",
    "                model=model,\n",
    "                scale=scale,\n",
    "                state_mean=state_mean,\n",
    "                state_std=state_std,\n",
    "                device=device,\n",
    "                target_return=TARGET_RETURN,\n",
    "            )\n",
    "\n",
    "            returns.append(ret)\n",
    "            lengths.append(ret)\n",
    "\n",
    "    return {\n",
    "        f'target_{TARGET_RETURN}_return_mean': np.mean(returns),\n",
    "        f'target_{TARGET_RETURN}_return_std': np.std(returns),\n",
    "        f'target_{TARGET_RETURN}_length_mean': np.mean(lengths),\n",
    "        f'target_{TARGET_RETURN}_length_std': np.std(lengths),\n",
    "    }\n",
    "\n",
    "class EvaluateCallback(TrainerCallback):\n",
    "    def on_epoch_end(self, args, state, control, **kwargs):\n",
    "        if int(state.epoch) % 500 == 0:\n",
    "            print('Epoch', state.epoch, 'eval:', evaluate_episodes(10, model))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "nNzzKWuuU9I4",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n"
     ]
    }
   ],
   "source": [
    "# these params more or less match the ones used by the original DT paper\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"output/\",\n",
    "    remove_unused_columns=False,\n",
    "    max_steps=50_000, # we only need about 50k steps to reach the highest score for DT and DM\n",
    "    logging_strategy='steps',\n",
    "    save_strategy='no',\n",
    "    logging_steps=5000 // 2,\n",
    "    warmup_steps=0,\n",
    "    per_device_train_batch_size=64,\n",
    "    learning_rate=1e-4,\n",
    "    weight_decay=1e-4,\n",
    "    #warmup_ratio=0.1,\n",
    "    optim=\"adamw_torch\",\n",
    "    dataloader_num_workers=16,\n",
    "    dataloader_persistent_workers=True,\n",
    "    max_grad_norm=0.25,\n",
    "    tf32=True,\n",
    "    bf16=True\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=dataset['train'],\n",
    "    data_collator=collator,\n",
    "    callbacks=[EvaluateCallback()],\n",
    ")\n",
    "\n",
    "#trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#trainer.save_model(f'trained_models/dt_{DATASET_NAME}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decision Mamba Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1287310\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='29505' max='50000' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [29505/50000 1:03:02 < 43:47, 7.80 it/s, Epoch 1844/3125]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>0.091300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5000</td>\n",
       "      <td>0.031600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7500</td>\n",
       "      <td>0.029500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10000</td>\n",
       "      <td>0.028300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>12500</td>\n",
       "      <td>0.027500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>15000</td>\n",
       "      <td>0.027100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>17500</td>\n",
       "      <td>0.026600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20000</td>\n",
       "      <td>0.026300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>22500</td>\n",
       "      <td>0.025900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>25000</td>\n",
       "      <td>0.025700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>27500</td>\n",
       "      <td>0.025700</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:36<00:00,  3.68s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 500.0 eval: {'target_12.0_return_mean': 11106.229210707726, 'target_12.0_return_std': 269.4713380655508, 'target_12.0_length_mean': 11106.229210707726, 'target_12.0_length_std': 269.4713380655508}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:35<00:00,  3.60s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1000.0 eval: {'target_12.0_return_mean': 11178.917994326002, 'target_12.0_return_std': 131.6579956579041, 'target_12.0_length_mean': 11178.917994326002, 'target_12.0_length_std': 131.6579956579041}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10/10 [00:36<00:00,  3.61s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1500.0 eval: {'target_12.0_return_mean': 11142.188534097539, 'target_12.0_return_std': 134.4963204642648, 'target_12.0_length_mean': 11142.188534097539, 'target_12.0_length_std': 134.4963204642648}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Exception ignored in: Exception ignored in: Exception ignored in: Exception ignored in: Exception ignored in: Exception ignored in: Exception ignored in: <bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>><bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Exception ignored in: <bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Exception ignored in: <bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Exception ignored in: <bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>><bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>><bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Exception ignored in: \n",
      "Exception ignored in: \n",
      "\n",
      "Exception ignored in: \n",
      "Exception ignored in: <bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>><bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>><bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>\n",
      "\n",
      "Traceback (most recent call last):\n",
      "\n",
      "<bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Traceback (most recent call last):\n",
      "<bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Traceback (most recent call last):\n",
      "\n",
      "Traceback (most recent call last):\n",
      "\n",
      "Traceback (most recent call last):\n",
      "\n",
      "Traceback (most recent call last):\n",
      "<bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Traceback (most recent call last):\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "\n",
      "\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "Traceback (most recent call last):\n",
      "<bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "Exception ignored in: \n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "    Traceback (most recent call last):\n",
      "      File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "      File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "\n",
      "            def _clean_thread_parent_frames(<bound method IPythonKernel._clean_thread_parent_frames of <ipykernel.ipkernel.IPythonKernel object at 0x7f0ee8835b40>>Traceback (most recent call last):\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "def _clean_thread_parent_frames(    def _clean_thread_parent_frames(    def _clean_thread_parent_frames(        Traceback (most recent call last):\n",
      "def _clean_thread_parent_frames(\n",
      "def _clean_thread_parent_frames(\n",
      "  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "\n",
      "    \n",
      "\n",
      "def _clean_thread_parent_frames(def _clean_thread_parent_frames(def _clean_thread_parent_frames(  File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "def _clean_thread_parent_frames(\n",
      "KeyboardInterrupt\n",
      "KeyboardInterruptTraceback (most recent call last):\n",
      "    def _clean_thread_parent_frames(KeyboardInterruptKeyboardInterrupt\n",
      "\n",
      "\n",
      "\n",
      "KeyboardInterruptKeyboardInterrupt: :       File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "def _clean_thread_parent_frames(\n",
      ": : KeyboardInterruptKeyboardInterruptKeyboardInterrupt: KeyboardInterrupt\n",
      "\n",
      "def _clean_thread_parent_frames(: \n",
      "    Traceback (most recent call last):\n",
      "KeyboardInterrupt: : \n",
      ": \n",
      "\n",
      ":   File \"/home/jvp/miniconda3/envs/mamba-ssm/lib/python3.10/site-packages/ipykernel/ipkernel.py\", line 770, in _clean_thread_parent_frames\n",
      "\n",
      "\n",
      "KeyboardInterrupt: def _clean_thread_parent_frames(\n",
      "\n",
      "\n",
      "\n",
      "    : \n",
      "\n",
      "KeyboardInterruptdef _clean_thread_parent_frames(\n",
      "KeyboardInterrupt: \n",
      ": \n",
      "KeyboardInterrupt\n",
      ": \n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[10], line 14\u001B[0m\n\u001B[1;32m      4\u001B[0m \u001B[38;5;28mprint\u001B[39m(num_params(model))\n\u001B[1;32m      6\u001B[0m trainer \u001B[38;5;241m=\u001B[39m Trainer(\n\u001B[1;32m      7\u001B[0m     model\u001B[38;5;241m=\u001B[39mmodel,\n\u001B[1;32m      8\u001B[0m     args\u001B[38;5;241m=\u001B[39mtraining_args,\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m     11\u001B[0m     callbacks\u001B[38;5;241m=\u001B[39m[EvaluateCallback()],\n\u001B[1;32m     12\u001B[0m )\n\u001B[0;32m---> 14\u001B[0m \u001B[43mtrainer\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mtrain\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     16\u001B[0m evaluate_episodes(\u001B[38;5;241m10\u001B[39m, model)\n",
      "File \u001B[0;32m~/.local/lib/python3.10/site-packages/transformers/trainer.py:1624\u001B[0m, in \u001B[0;36mTrainer.train\u001B[0;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001B[0m\n\u001B[1;32m   1622\u001B[0m         hf_hub_utils\u001B[38;5;241m.\u001B[39menable_progress_bars()\n\u001B[1;32m   1623\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m-> 1624\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43minner_training_loop\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m   1625\u001B[0m \u001B[43m        \u001B[49m\u001B[43margs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m   1626\u001B[0m \u001B[43m        \u001B[49m\u001B[43mresume_from_checkpoint\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mresume_from_checkpoint\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m   1627\u001B[0m \u001B[43m        \u001B[49m\u001B[43mtrial\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mtrial\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m   1628\u001B[0m \u001B[43m        \u001B[49m\u001B[43mignore_keys_for_eval\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mignore_keys_for_eval\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m   1629\u001B[0m \u001B[43m    \u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/.local/lib/python3.10/site-packages/transformers/trainer.py:1928\u001B[0m, in \u001B[0;36mTrainer._inner_training_loop\u001B[0;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001B[0m\n\u001B[1;32m   1925\u001B[0m     rng_to_sync \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mTrue\u001B[39;00m\n\u001B[1;32m   1927\u001B[0m step \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m1\u001B[39m\n\u001B[0;32m-> 1928\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m step, inputs \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28menumerate\u001B[39m(epoch_iterator):\n\u001B[1;32m   1929\u001B[0m     total_batched_samples \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[1;32m   1931\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39margs\u001B[38;5;241m.\u001B[39minclude_num_input_tokens_seen:\n",
      "File \u001B[0;32m~/.local/lib/python3.10/site-packages/accelerate/data_loader.py:452\u001B[0m, in \u001B[0;36mDataLoaderShard.__iter__\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    450\u001B[0m \u001B[38;5;66;03m# We iterate one batch ahead to check when we are at the end\u001B[39;00m\n\u001B[1;32m    451\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m--> 452\u001B[0m     current_batch \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mnext\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mdataloader_iter\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    453\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mStopIteration\u001B[39;00m:\n\u001B[1;32m    454\u001B[0m     \u001B[38;5;28;01myield\u001B[39;00m\n",
      "File \u001B[0;32m~/.local/lib/python3.10/site-packages/torch/utils/data/dataloader.py:631\u001B[0m, in \u001B[0;36m_BaseDataLoaderIter.__next__\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    628\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_sampler_iter \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[1;32m    629\u001B[0m     \u001B[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001B[39;00m\n\u001B[1;32m    630\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_reset()  \u001B[38;5;66;03m# type: ignore[call-arg]\u001B[39;00m\n\u001B[0;32m--> 631\u001B[0m data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_next_data\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    632\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_num_yielded \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[1;32m    633\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_dataset_kind \u001B[38;5;241m==\u001B[39m _DatasetKind\u001B[38;5;241m.\u001B[39mIterable \u001B[38;5;129;01mand\u001B[39;00m \\\n\u001B[1;32m    634\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_IterableDataset_len_called \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m \\\n\u001B[1;32m    635\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_num_yielded \u001B[38;5;241m>\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_IterableDataset_len_called:\n",
      "File \u001B[0;32m~/.local/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1329\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._next_data\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m   1326\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_process_data(data)\n\u001B[1;32m   1328\u001B[0m \u001B[38;5;28;01massert\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_shutdown \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_tasks_outstanding \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m0\u001B[39m\n\u001B[0;32m-> 1329\u001B[0m idx, data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_get_data\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1330\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_tasks_outstanding \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[1;32m   1331\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_dataset_kind \u001B[38;5;241m==\u001B[39m _DatasetKind\u001B[38;5;241m.\u001B[39mIterable:\n\u001B[1;32m   1332\u001B[0m     \u001B[38;5;66;03m# Check for _IterableDatasetStopIteration\u001B[39;00m\n",
      "File \u001B[0;32m~/.local/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1285\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._get_data\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m   1283\u001B[0m \u001B[38;5;28;01melif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_pin_memory:\n\u001B[1;32m   1284\u001B[0m     \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_pin_memory_thread\u001B[38;5;241m.\u001B[39mis_alive():\n\u001B[0;32m-> 1285\u001B[0m         success, data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_try_get_data\u001B[49m\u001B[43m(\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1286\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m success:\n\u001B[1;32m   1287\u001B[0m             \u001B[38;5;28;01mreturn\u001B[39;00m data\n",
      "File \u001B[0;32m~/.local/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1133\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001B[0;34m(self, timeout)\u001B[0m\n\u001B[1;32m   1120\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21m_try_get_data\u001B[39m(\u001B[38;5;28mself\u001B[39m, timeout\u001B[38;5;241m=\u001B[39m_utils\u001B[38;5;241m.\u001B[39mMP_STATUS_CHECK_INTERVAL):\n\u001B[1;32m   1121\u001B[0m     \u001B[38;5;66;03m# Tries to fetch data from `self._data_queue` once for a given timeout.\u001B[39;00m\n\u001B[1;32m   1122\u001B[0m     \u001B[38;5;66;03m# This can also be used as inner loop of fetching without timeout, with\u001B[39;00m\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m   1130\u001B[0m     \u001B[38;5;66;03m# Returns a 2-tuple:\u001B[39;00m\n\u001B[1;32m   1131\u001B[0m     \u001B[38;5;66;03m#   (bool: whether successfully get data, any: data if successful else None)\u001B[39;00m\n\u001B[1;32m   1132\u001B[0m     \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[0;32m-> 1133\u001B[0m         data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_data_queue\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mget\u001B[49m\u001B[43m(\u001B[49m\u001B[43mtimeout\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mtimeout\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1134\u001B[0m         \u001B[38;5;28;01mreturn\u001B[39;00m (\u001B[38;5;28;01mTrue\u001B[39;00m, data)\n\u001B[1;32m   1135\u001B[0m     \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mException\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:\n\u001B[1;32m   1136\u001B[0m         \u001B[38;5;66;03m# At timeout and error, we manually check whether any worker has\u001B[39;00m\n\u001B[1;32m   1137\u001B[0m         \u001B[38;5;66;03m# failed. Note that this is the only mechanism for Windows to detect\u001B[39;00m\n\u001B[1;32m   1138\u001B[0m         \u001B[38;5;66;03m# worker failures.\u001B[39;00m\n",
      "File \u001B[0;32m~/miniconda3/envs/mamba-ssm/lib/python3.10/queue.py:180\u001B[0m, in \u001B[0;36mQueue.get\u001B[0;34m(self, block, timeout)\u001B[0m\n\u001B[1;32m    178\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m remaining \u001B[38;5;241m<\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m0.0\u001B[39m:\n\u001B[1;32m    179\u001B[0m             \u001B[38;5;28;01mraise\u001B[39;00m Empty\n\u001B[0;32m--> 180\u001B[0m         \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mnot_empty\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mwait\u001B[49m\u001B[43m(\u001B[49m\u001B[43mremaining\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    181\u001B[0m item \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_get()\n\u001B[1;32m    182\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mnot_full\u001B[38;5;241m.\u001B[39mnotify()\n",
      "File \u001B[0;32m~/miniconda3/envs/mamba-ssm/lib/python3.10/threading.py:324\u001B[0m, in \u001B[0;36mCondition.wait\u001B[0;34m(self, timeout)\u001B[0m\n\u001B[1;32m    322\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m    323\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m timeout \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m0\u001B[39m:\n\u001B[0;32m--> 324\u001B[0m         gotit \u001B[38;5;241m=\u001B[39m \u001B[43mwaiter\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43macquire\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43;01mTrue\u001B[39;49;00m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtimeout\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m    325\u001B[0m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[1;32m    326\u001B[0m         gotit \u001B[38;5;241m=\u001B[39m waiter\u001B[38;5;241m.\u001B[39macquire(\u001B[38;5;28;01mFalse\u001B[39;00m)\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "config = DecisionTransformerConfig(state_dim=collator.state_dim, act_dim=collator.act_dim)\n",
    "model = TrainableDM(config)\n",
    "\n",
    "print(num_params(model))\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=dataset['train'],\n",
    "    data_collator=collator,\n",
    "    callbacks=[EvaluateCallback()],\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "\n",
    "evaluate_episodes(10, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sNaj6bOkp3bt"
   },
   "source": [
    "### Step 7: Visualize the performance of the agent\n",
    "\n",
    "With mujoco_py, it'll take a little while to compile the first "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(f'trained_models/dm_{DATASET_NAME}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "private_outputs": true,
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}