CSE571_HW3.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lwxyFyqU7PM-"
      },
      "outputs": [],
      "source": [
        "# upload files listed below from the folder (or you can simply upload all files)\n",
        "# The `files` tab will be like\n",
        "# - sample_data (this one is default in colab)\n",
        "# - utils.py\n",
        "# - policy.py\n",
        "# - evaluate.py\n",
        "# - pytorch_utils.py\n",
        "# - expert_data.pkl\n",
        "# - expert_policy.pkl\n",
        "\n",
        "# !NOTE!: You need to copy your implementation of bc, dagger and policy_gradient \n",
        "# in the notebook here back to the python script when you submiting your code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6eG6uXB6BCb7",
        "outputId": "9c3cd241-d785-41b4-8fc3-ca80ff6e5722"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "libglew-dev is already the newest version (2.1.0-4).\n",
            "libgl1-mesa-dev is already the newest version (21.2.6-0ubuntu0.1~20.04.2).\n",
            "libgl1-mesa-glx is already the newest version (21.2.6-0ubuntu0.1~20.04.2).\n",
            "libosmesa6-dev is already the newest version (21.2.6-0ubuntu0.1~20.04.2).\n",
            "software-properties-common is already the newest version (0.99.9.11).\n",
            "0 upgraded, 0 newly installed, 0 to remove and 27 not upgraded.\n",
            "Reading package lists... Done\n",
            "Building dependency tree       \n",
            "Reading state information... Done\n",
            "patchelf is already the newest version (0.10-2build1).\n",
            "0 upgraded, 0 newly installed, 0 to remove and 27 not upgraded.\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: free-mujoco-py in /usr/local/lib/python3.10/dist-packages (2.1.6)\n",
            "Requirement already satisfied: Cython<0.30.0,>=0.29.24 in /usr/local/lib/python3.10/dist-packages (from free-mujoco-py) (0.29.34)\n",
            "Requirement already satisfied: cffi<2.0.0,>=1.15.0 in /usr/local/lib/python3.10/dist-packages (from free-mujoco-py) (1.15.1)\n",
            "Requirement already satisfied: fasteners==0.15 in /usr/local/lib/python3.10/dist-packages (from free-mujoco-py) (0.15)\n",
            "Requirement already satisfied: glfw<2.0.0,>=1.4.0 in /usr/local/lib/python3.10/dist-packages (from free-mujoco-py) (1.12.0)\n",
            "Requirement already satisfied: imageio<3.0.0,>=2.9.0 in /usr/local/lib/python3.10/dist-packages (from free-mujoco-py) (2.25.1)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.21.3 in /usr/local/lib/python3.10/dist-packages (from free-mujoco-py) (1.22.4)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.10/dist-packages (from fasteners==0.15->free-mujoco-py) (1.16.0)\n",
            "Requirement already satisfied: monotonic>=0.1 in /usr/local/lib/python3.10/dist-packages (from fasteners==0.15->free-mujoco-py) (1.6)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.10/dist-packages (from cffi<2.0.0,>=1.15.0->free-mujoco-py) (2.21)\n",
            "Requirement already satisfied: pillow>=8.3.2 in /usr/local/lib/python3.10/dist-packages (from imageio<3.0.0,>=2.9.0->free-mujoco-py) (8.4.0)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: gym==0.22.0 in /usr/local/lib/python3.10/dist-packages (0.22.0)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.10/dist-packages (from gym==0.22.0) (1.22.4)\n",
            "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.10/dist-packages (from gym==0.22.0) (2.2.1)\n",
            "Requirement already satisfied: gym-notices>=0.0.4 in /usr/local/lib/python3.10/dist-packages (from gym==0.22.0) (0.0.8)\n"
          ]
        }
      ],
      "source": [
        "!apt-get install -y \\\n",
        "   libgl1-mesa-dev \\\n",
        "   libgl1-mesa-glx \\\n",
        "   libglew-dev \\\n",
        "   libosmesa6-dev \\\n",
        "   software-properties-common\n",
        "\n",
        "!apt-get install -y patchelf \n",
        "\n",
        "!pip install free-mujoco-py\n",
        "!pip install gym==0.22.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7tqEtT1mAhYO"
      },
      "outputs": [],
      "source": [
        "#Include this at the top of your colab code\n",
        "import os\n",
        "if not os.path.exists('.mujoco_setup_complete'):\n",
        "  # Get the prereqs\n",
        "  !apt-get -qq update\n",
        "  !apt-get -qq install -y libosmesa6-dev libgl1-mesa-glx libglfw3 libgl1-mesa-dev libglew-dev patchelf\n",
        "  # Get Mujoco\n",
        "  !mkdir ~/.mujoco\n",
        "  !wget -q https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz -O mujoco.tar.gz\n",
        "  !tar -zxf mujoco.tar.gz -C \"$HOME/.mujoco\"\n",
        "  !rm mujoco.tar.gz\n",
        "  # Add it to the actively loaded path and the bashrc path (these only do so much)\n",
        "  !echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco210/bin' >> ~/.bashrc \n",
        "  !echo 'export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libGLEW.so' >> ~/.bashrc \n",
        "  # THE ANNOYING ONE, FORCE IT INTO LDCONFIG SO WE ACTUALLY GET ACCESS TO IT THIS SESSION\n",
        "  !echo \"/root/.mujoco/mujoco210/bin\" > /etc/ld.so.conf.d/mujoco_ld_lib_path.conf\n",
        "  !ldconfig\n",
        "  # Install Mujoco-py\n",
        "  !pip3 install -U 'mujoco-py<2.2,>=2.1'\n",
        "  # run once\n",
        "  !touch .mujoco_setup_complete\n",
        "\n",
        "try:\n",
        "  if _mujoco_run_once:\n",
        "    pass\n",
        "except NameError:\n",
        "  _mujoco_run_once = False\n",
        "if not _mujoco_run_once:\n",
        "  # Add it to the actively loaded path and the bashrc path (these only do so much)\n",
        "  try:\n",
        "    os.environ['LD_LIBRARY_PATH']=os.environ['LD_LIBRARY_PATH'] + ':/root/.mujoco/mujoco210/bin'\n",
        "  except KeyError:\n",
        "    os.environ['LD_LIBRARY_PATH']='/root/.mujoco/mujoco210/bin'\n",
        "  try:\n",
        "    os.environ['LD_PRELOAD']=os.environ['LD_PRELOAD'] + ':/usr/lib/x86_64-linux-gnu/libGLEW.so'\n",
        "  except KeyError:\n",
        "    os.environ['LD_PRELOAD']='/usr/lib/x86_64-linux-gnu/libGLEW.so'\n",
        "  # presetup so we don't see output on first env initialization\n",
        "  import mujoco_py\n",
        "  _mujoco_run_once = True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GaJ_54U1BRu9",
        "outputId": "f0f7266a-b53e-493f-c715-02ed30dc877f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(11,)\n"
          ]
        }
      ],
      "source": [
        "# Test whether mujoco and gym is correctly installed\n",
        "# You are expected to get an output of `(11, )`\n",
        "import gym\n",
        "import mujoco_py\n",
        "env = gym.make(\"Reacher-v2\")\n",
        "print(env.reset().shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.optim as optim\n",
        "import numpy as np\n",
        "from utils import log_density, rollout\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R7iC3MtQDIUm"
      },
      "outputs": [],
      "source": [
        "# NOTE: copy your implementation back to bc.py when you submit the code\n",
        "def simulate_policy_bc(env, policy, expert_data, num_epochs=500, episode_length=50, \n",
        "                       batch_size=32):\n",
        "    \n",
        "    # Hint: Just flatten your expert dataset and use standard pytorch supervised learning code to train the policy. \n",
        "    optimizer = optim.Adam(list(policy.parameters()))\n",
        "    idxs = np.array(range(len(expert_data)))\n",
        "    num_batches = len(idxs)*episode_length // batch_size\n",
        "    losses = []\n",
        "    for epoch in range(num_epochs): \n",
        "        ## TODO Students\n",
        "        np.random.shuffle(idxs)\n",
        "        running_loss = 0.0\n",
        "        for i in range(num_batches):\n",
        "            optimizer.zero_grad()\n",
        "            # TODO start: Fill in your behavior cloning implementation here\n",
        "            \n",
        "            # TODO end\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            running_loss += loss.item()\n",
        "        if epoch % 50 == 0:\n",
        "            print('[%d] loss: %.8f' %\n",
        "                (epoch, running_loss / 10.))\n",
        "        losses.append(loss.item())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q4pUKeaACk_k"
      },
      "outputs": [],
      "source": [
        "# NOTE: copy your implementation back to dagger.py when you submit the code\n",
        "def simulate_policy_dagger(env, policy, expert_paths, expert_policy=None, num_epochs=500, episode_length=50,\n",
        "                            batch_size=32, num_dagger_iters=10, num_trajs_per_dagger=10):\n",
        "    \n",
        "    # TODO: Fill in your dagger implementation here. Make sure to plot your intermediate returns\n",
        "    \n",
        "    # Hint: Loop through num_dagger_iters iterations, at each iteration train a policy on the current dataset.\n",
        "    # Then rollout the policy, use relabel_action to relabel the actions along the trajectory with \"expert_policy\" and then add this to current dataset\n",
        "    # Repeat this so the dataset grows with states drawn from the policy, and relabeled actions using the expert.\n",
        "    # Optimizer code\n",
        "    optimizer = optim.Adam(list(policy.parameters()))\n",
        "    losses = []\n",
        "    returns = []\n",
        "\n",
        "    trajs = expert_paths\n",
        "    # Dagger iterations\n",
        "    for dagger_itr in range(num_dagger_iters):\n",
        "        idxs = np.array(range(len(trajs)))\n",
        "        num_batches = len(idxs)*episode_length // batch_size\n",
        "        losses = []\n",
        "        # Train the model with Adam\n",
        "        for epoch in range(num_epochs):\n",
        "            running_loss = 0.0\n",
        "            for i in range(num_batches):\n",
        "                optimizer.zero_grad()\n",
        "                # TODO start: Fill in your behavior cloning implementation here\n",
        "                \n",
        "                # TODO end\n",
        "                loss.backward()\n",
        "                optimizer.step()\n",
        "\n",
        "                # print statistics\n",
        "                running_loss += loss.item()\n",
        "            # print('[%d, %5d] loss: %.8f' %(epoch + 1, i + 1, running_loss))\n",
        "            losses.append(loss.item())\n",
        "\n",
        "        # Collecting more data for dagger\n",
        "        trajs_recent = []\n",
        "        for k in range(num_trajs_per_dagger):\n",
        "            env.reset()\n",
        "            # TODO start: Rollout the policy on the environment to collect more data, relabel them, add them into trajs_recent\n",
        "            \n",
        "            # TODO end\n",
        "\n",
        "        trajs += trajs_recent\n",
        "        mean_return = np.mean(np.array([traj['rewards'].sum() for traj in trajs_recent]))\n",
        "        print(\"Average DAgger return is \" + str(mean_return))\n",
        "        returns.append(mean_return)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9mr2-0nOCwo8"
      },
      "outputs": [],
      "source": [
        "# NOTE: copy your implementation back to policy_gradient.py when you submit the code\n",
        "def train_model(policy, baseline, trajs, policy_optim, baseline_optim, gamma=0.99, baseline_train_batch_size=64, baseline_num_epochs=5):\n",
        "    # Fill in your policy gradient implementation here\n",
        "\n",
        "    # TODO: Compute the returns on the current batch of trajectories\n",
        "    # Hint: Go through all the trajectories in trajs and compute their return to go: discounted sum of rewards from that timestep to the end. \n",
        "    # Hint: This is easy to do if you go backwards in time and sum up the reward as a running sum. \n",
        "    # Hint: Remember that return to go is return = r[t] + gamma*r[t+1] + gamma^2*r[t+2] + ...\n",
        "    states_all = []\n",
        "    actions_all = []\n",
        "    returns_all = []\n",
        "    for traj in trajs:\n",
        "        states_singletraj = traj['observations']\n",
        "        actions_singletraj = traj['actions']\n",
        "        rewards_singletraj = traj['rewards']\n",
        "        returns_singletraj = np.zeros_like(rewards_singletraj)\n",
        "        # TODO start\n",
        "            \n",
        "        # TODO end\n",
        "        states_all.append(states_singletraj)\n",
        "        actions_all.append(actions_singletraj)\n",
        "        returns_all.append(returns_singletraj)\n",
        "    states = np.concatenate(states_all)\n",
        "    actions = np.concatenate(actions_all)\n",
        "    returns = np.concatenate(returns_all)\n",
        "    \n",
        "    # TODO: Normalize the returns by subtracting mean and dividing by std\n",
        "    # Hint: Just do return - return.mean()/ (return.std() + EPS), where EPS is a small constant for numerics\n",
        "    # TODO start\n",
        "    \n",
        "    # TODO end\n",
        "    \n",
        "    # TODO: Train baseline by regressing onto returns\n",
        "    # Hint: Regress the baseline from each state onto the above computed return to go. You can use similar code to behavior cloning to do so. \n",
        "    # Hint: Iterate for baseline_num_epochs with batch size = baseline_train_batch_size\n",
        "    criterion = torch.nn.MSELoss()\n",
        "    n = len(states)\n",
        "    arr = np.arange(n)\n",
        "    for epoch in range(baseline_num_epochs):\n",
        "        np.random.shuffle(arr)\n",
        "        for i in range(n // baseline_train_batch_size):\n",
        "            batch_index = arr[baseline_train_batch_size * i: baseline_train_batch_size * (i + 1)]\n",
        "            batch_index = torch.LongTensor(batch_index).to(device)\n",
        "            # TODO start         \n",
        "            \n",
        "            # TODO end\n",
        "            baseline_optim.zero_grad()\n",
        "            loss.backward()\n",
        "            baseline_optim.step()\n",
        "            \n",
        "    # TODO: Train policy by optimizing surrogate objective: -log prob * (return - baseline)\n",
        "    # Hint: Policy gradient is given by: \\grad log prob(a|s)* (return - baseline)\n",
        "    # Hint: Return is computed above, you can computer log_probs using the log_density function imported. \n",
        "    # Hint: You can predict what the baseline outputs for every state.  \n",
        "    # Hint: Then simply compute the surrogate objective by taking the objective as -log prob * (return - baseline)\n",
        "    # Hint: You can then use standard pytorch machinery to take *one* gradient step on the policy\n",
        "    mu, std, logstd = policy(torch.Tensor(states).to(device))\n",
        "    log_policy = log_density(torch.Tensor(actions).to(device), mu, std, logstd)\n",
        "    baseline_pred = baseline(torch.from_numpy(states).float().to(device))\n",
        "    # TODO start\n",
        "    \n",
        "    # TODO end\n",
        "    \n",
        "    policy_optim.zero_grad()\n",
        "    loss.backward()\n",
        "    policy_optim.step()\n",
        "    \n",
        "    del states, actions, returns, states_all, actions_all, returns_all\n",
        "\n",
        "# Training loop for policy gradient\n",
        "def simulate_policy_pg(env, policy, baseline, num_epochs=20000, max_path_length=200, pg_batch_size=100, \n",
        "                        gamma=0.99, baseline_train_batch_size=64, baseline_num_epochs=5, print_freq=10, render=False):\n",
        "    policy_optim = optim.Adam(policy.parameters())\n",
        "    baseline_optim = optim.Adam(baseline.parameters())\n",
        "\n",
        "    for iter_num in range(num_epochs):\n",
        "        sample_trajs = []\n",
        "\n",
        "        # Sampling trajectories\n",
        "        for it in range(pg_batch_size):\n",
        "            sample_traj = rollout(env=env, agent=policy, episode_length=max_path_length, agent_name='pg', render=render)\n",
        "            sample_trajs.append(sample_traj)\n",
        "        \n",
        "        # Logging returns occasionally\n",
        "        if iter_num % print_freq == 0:\n",
        "            rewards_np = np.mean(np.asarray([traj['rewards'].sum() for traj in sample_trajs]))\n",
        "            path_length = np.max(np.asarray([traj['rewards'].shape[0] for traj in sample_trajs]))\n",
        "            print(\"Episode: {}, reward: {}, max path length: {}\".format(iter_num, rewards_np, path_length))\n",
        "\n",
        "        # Training model\n",
        "        train_model(policy, baseline, sample_trajs, policy_optim, baseline_optim, gamma=gamma, \n",
        "                    baseline_train_batch_size=baseline_train_batch_size, baseline_num_epochs=baseline_num_epochs)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pQNm7Fj1DqT6",
        "outputId": "86a8a8f9-5e3e-48d5-cf18-ec1b82c2b792"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/ipykernel/ipkernel.py:283: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
            "  and should_run_async(code)\n"
          ]
        }
      ],
      "source": [
        "\n",
        "class Args:\n",
        "    def __init__(self, task, test, render=False):\n",
        "        self.task = task # 'behavior_cloning', 'dagger' or 'policy_gradient'\n",
        "        self.test = test # whether test only\n",
        "        self.render = render # whether to render during test\n",
        "args = Args('policy_gradient', False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 607
        },
        "id": "D3T_K8lZC4jt",
        "outputId": "0d34dc47-2251-4ad1-e829-5b2ccab7bd3b"
      },
      "outputs": [],
      "source": [
        "import gym\n",
        "import numpy as np\n",
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "import torch\n",
        "import argparse\n",
        "import torch.nn as nn\n",
        "\n",
        "from utils import BCPolicy, generate_paths, get_expert_data, RLPolicy, RLBaseline\n",
        "from policy import MakeDeterministic\n",
        "import pytorch_utils as ptu\n",
        "from evaluate import evaluate\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print('using device', device)\n",
        "\n",
        "torch.manual_seed(0)\n",
        "import random\n",
        "random.seed(0)\n",
        "np.random.seed(0)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "\n",
        "    if args.task == 'policy_gradient':\n",
        "        # Define environment\n",
        "        env = gym.make(\"InvertedPendulum-v2\")\n",
        "\n",
        "        # Define policy and value function\n",
        "        hidden_dim_pol = 64\n",
        "        hidden_depth_pol = 2\n",
        "        hidden_dim_baseline = 64\n",
        "        hidden_depth_baseline = 2\n",
        "        policy = RLPolicy(env.observation_space.shape[0], env.action_space.shape[0], hidden_dim=hidden_dim_pol, hidden_depth=hidden_depth_pol)\n",
        "        baseline = RLBaseline(env.observation_space.shape[0], hidden_dim=hidden_dim_baseline, hidden_depth=hidden_depth_baseline)\n",
        "        policy.to(device)\n",
        "        baseline.to(device)\n",
        "\n",
        "        # Training hyperparameters\n",
        "        num_epochs=100\n",
        "        max_path_length=200\n",
        "        pg_batch_size=100\n",
        "        gamma=0.99\n",
        "        baseline_train_batch_size=64\n",
        "        baseline_num_epochs=5\n",
        "        print_freq=10\n",
        "\n",
        "        if not args.test:\n",
        "            # Train policy gradient\n",
        "            simulate_policy_pg(env, policy, baseline, num_epochs=num_epochs, max_path_length=max_path_length, pg_batch_size=pg_batch_size, \n",
        "                            gamma=gamma, baseline_train_batch_size=baseline_train_batch_size, baseline_num_epochs=baseline_num_epochs, print_freq=print_freq, render=False)\n",
        "            torch.save(policy.state_dict(), 'pg_final.pth')\n",
        "        else:\n",
        "            policy.load_state_dict(torch.load(f'pg_final.pth'))\n",
        "        evaluate(env, policy, 'pg', num_validation_runs=100, episode_length=max_path_length, render=args.render)\n",
        "\n",
        "    if args.task == 'behavior_cloning':\n",
        "        # Define environment\n",
        "        env = gym.make(\"Reacher-v2\")\n",
        "\n",
        "        # Define policy\n",
        "        hidden_dim = 128\n",
        "        hidden_depth = 2\n",
        "        obs_size = env.observation_space.shape[0]\n",
        "        ac_size = env.action_space.shape[0]\n",
        "        print(obs_size, ac_size, hidden_dim, hidden_depth)\n",
        "        policy = BCPolicy(obs_size, ac_size, hidden_dim=hidden_dim, hidden_depth=hidden_depth) # 10 dimensional latent\n",
        "        policy.to(device)\n",
        "\n",
        "        # Get the expert data\n",
        "        file_path = 'expert_data.pkl'\n",
        "        expert_data = get_expert_data(file_path)\n",
        "        # print(expert_data)\n",
        "        # Training hyperparameters\n",
        "        episode_length = 50\n",
        "        num_epochs = 500\n",
        "        batch_size = 32\n",
        "        num_validation_runs = 100\n",
        "\n",
        "        if not args.test:\n",
        "        # Train behavior cloning\n",
        "            simulate_policy_bc(env, policy, expert_data, num_epochs=num_epochs, episode_length=episode_length,\n",
        "                            batch_size=batch_size)\n",
        "            torch.save(policy.state_dict(), 'bc_final.pth')\n",
        "        else:\n",
        "            policy.load_state_dict(torch.load(f'bc_final.pth'))\n",
        "        evaluate(env, policy, 'bc', num_validation_runs=num_validation_runs, episode_length=episode_length, render=args.render)\n",
        "\n",
        "    if args.task == 'dagger':\n",
        "        # Define environment\n",
        "        env = gym.make(\"Reacher-v2\")\n",
        "\n",
        "        # Policy\n",
        "        hidden_dim = 1000\n",
        "        hidden_depth = 3\n",
        "        obs_size = env.observation_space.shape[0]\n",
        "        ac_size = env.action_space.shape[0]\n",
        "        policy = BCPolicy(obs_size, ac_size, hidden_dim=hidden_dim, hidden_depth=hidden_depth) # 10 dimensional latent\n",
        "        policy.to(device)\n",
        "\n",
        "        # Get the expert data\n",
        "        file_path = 'expert_data.pkl'\n",
        "        expert_data = get_expert_data(file_path)\n",
        "\n",
        "        # Load interactive expert\n",
        "        expert_policy = torch.load('expert_policy.pkl', map_location=torch.device(device))\n",
        "        print(\"Expert policy loaded\")\n",
        "        expert_policy.to(device)\n",
        "        ptu.set_gpu_mode(True)\n",
        "\n",
        "        # Training hyperparameters\n",
        "        episode_length = 50\n",
        "        num_epochs = 400\n",
        "        batch_size = 32\n",
        "        num_dagger_iters = 10\n",
        "        num_trajs_per_dagger = 10\n",
        "        num_validation_runs = 100\n",
        "\n",
        "        if not args.test:\n",
        "        # Train DAgger\n",
        "            simulate_policy_dagger(env, policy, expert_data, expert_policy, num_epochs=num_epochs, episode_length=episode_length,\n",
        "                               batch_size=batch_size, num_dagger_iters=num_dagger_iters, num_trajs_per_dagger=num_trajs_per_dagger)\n",
        "            torch.save(policy.state_dict(), 'dagger_final.pth')\n",
        "        else:\n",
        "            policy.load_state_dict(torch.load(f'dagger_final.pth'))\n",
        "        evaluate(env, policy, 'dagger', num_validation_runs=num_validation_runs, episode_length=episode_length, render=args.render)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}