DQN_generate_dataset.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Introduction\n",
        "\n",
        "This notebook generates a CarRacing-v2 dataset for offline reinforcement learning by using a pretrained DQN for environment exploration, with configurable episode count and length, and stores it in HDF5 on Google Cloud Storage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpjcCHxkeLKO"
      },
      "source": [
        "### Install initial environment in Google Colab"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GRx1SfZkeLKQ",
        "outputId": "7cb88885-ea76-40bf-db14-2c9cdcb1373d"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "import os\n",
        "\n",
        "if 'google.colab' in sys.modules:\n",
        "  if not os.path.exists('/content/.already_installed'):\n",
        "    !git clone https://github.com/YakivGalkin/cnn_decision_transformer\n",
        "    !apt-get install -y swig\n",
        "    !pip install -r cnn_decision_transformer/requirements.txt\n",
        "    with open('/content/.already_installed', 'w') as f:\n",
        "        f.write('done')\n",
        "  %cd /content/cnn_decision_transformer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tyFxDs1beLKR"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "99O1Dp3weLKR",
        "outputId": "caaf2cfe-d075-495a-857b-fa23d9426e62"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/jacob/.pyenv/versions/3.10.12/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "from dataclasses import dataclass\n",
        "import gymnasium as gym\n",
        "import numpy as np\n",
        "import torch\n",
        "from stable_baselines3 import DQN\n",
        "from stable_baselines3.common.torch_layers import NatureCNN\n",
        "from stable_baselines3.common.vec_env import DummyVecEnv, VecTransposeImage\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F4p2pRjAeLKS"
      },
      "source": [
        "### Generate Data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load pretrained DQN Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import utils.storage as storage\n",
        "loaded_model = storage.get_pretrained_model(DQN,  'dql_car_racing.pt')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Iteract with environmenr, build trajectories"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UhM2DabQeLKT",
        "outputId": "21e5d08d-53d4-4b2c-aa0d-9ff8d77d431b"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/jacob/.pyenv/versions/3.10.12/lib/python3.10/site-packages/stable_baselines3/common/save_util.py:166: UserWarning: Could not deserialize object lr_schedule. Consider using `custom_objects` argument to replace this object.\n",
            "Exception: Can't get attribute '_make_function' on <module 'cloudpickle.cloudpickle' from '/Users/jacob/.pyenv/versions/3.10.12/lib/python3.10/site-packages/cloudpickle/cloudpickle.py'>\n",
            "  warnings.warn(\n",
            "/Users/jacob/.pyenv/versions/3.10.12/lib/python3.10/site-packages/stable_baselines3/common/save_util.py:166: UserWarning: Could not deserialize object exploration_schedule. Consider using `custom_objects` argument to replace this object.\n",
            "Exception: Can't get attribute '_make_function' on <module 'cloudpickle.cloudpickle' from '/Users/jacob/.pyenv/versions/3.10.12/lib/python3.10/site-packages/cloudpickle/cloudpickle.py'>\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Episode: 0 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 22.25806451612903 episodes steps: 100\n",
            "Episode: 1 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 16.49006622516558 episodes steps: 100\n",
            "Episode: 2 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 6.835016835016857 episodes steps: 100\n",
            "Episode: 3 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 9.455252918287968 episodes steps: 100\n",
            "Episode: 4 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 23.003300330033 episodes steps: 100\n",
            "Episode: 5 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 14.734982332155512 episodes steps: 100\n",
            "Episode: 6 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 17.777777777777786 episodes steps: 100\n",
            "Episode: 7 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 17.30375426621162 episodes steps: 100\n",
            "Episode: 8 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 18.846153846153847 episodes steps: 100\n",
            "Episode: 9 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 14.390243902439058 episodes steps: 100\n",
            "Episode: 10 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 17.777777777777793 episodes steps: 100\n",
            "Episode: 11 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 26.231884057970994 episodes steps: 100\n",
            "Episode: 12 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 15.925925925925961 episodes steps: 100\n",
            "Episode: 13 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 10.066889632107058 episodes steps: 100\n",
            "Episode: 14 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 12.222222222222253 episodes steps: 100\n",
            "Episode: 15 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 21.14186851211071 episodes steps: 100\n",
            "Episode: 16 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 15.641025641025667 episodes steps: 100\n",
            "Episode: 17 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 17.700831024930764 episodes steps: 100\n",
            "Episode: 18 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 11.897810218978139 episodes steps: 100\n",
            "Episode: 19 of 20:\n",
            "....................................................................................................\n",
            "Total reward: 17.86377708978329 episodes steps: 100\n"
          ]
        }
      ],
      "source": [
        "from visual_decision_transformer.cnn_feature_extractor import prepare_env_observation_array\n",
        "\n",
        "\n",
        "env =  gym.make('CarRacing-v2', continuous=False) \n",
        "\n",
        "MAX_EPISODE_STEPS = 1000 \n",
        "NUM_EPISODES = 150\n",
        "RTG_GAMMA = 1.0 # \"Return to GO\" with the default GAMMA is stored in the initial dataset. \n",
        "CUT_IF_NO_REWARD = 50\n",
        "\n",
        "\n",
        "features = {\n",
        "    \"observations\": [],\n",
        "    \"actions\": [],\n",
        "    \"rewards\": [],\n",
        "    \"dones\": [],\n",
        "    \"rtg\": []\n",
        "}\n",
        "\n",
        "for episode in range(NUM_EPISODES):\n",
        "    print(f\"Episode: {episode} of {NUM_EPISODES}:\" )\n",
        "    [obs, _] = env.reset()\n",
        "    done = False\n",
        "\n",
        "    o, a, r, d, g = [], [], [], [], []\n",
        "    total_reward = 0\n",
        "    step_index = 0\n",
        "    while not done:\n",
        "        step_index +=  1\n",
        "\n",
        "        action, _states = loaded_model.predict(obs,deterministic=True)\n",
        "        new_obs, reward, done, t, i = env.step(action)\n",
        "        total_reward = total_reward + reward\n",
        "        oarr = prepare_env_observation_array(obs)\n",
        "        o.append(oarr.flatten())\n",
        "        a.append(action.item())\n",
        "        r.append(reward)\n",
        "        d.append(done)\n",
        "        obs = new_obs\n",
        "        print(\".\", end=\"\")\n",
        "\n",
        "        # any positive reward within last 50 steps?\n",
        "        if step_index > MAX_EPISODE_STEPS or (len(r) > CUT_IF_NO_REWARD and max(r[-CUT_IF_NO_REWARD:]) <= 0):\n",
        "            break\n",
        "\n",
        "    fin_positive = len(r)-1\n",
        "    for fin_positive in range (len(r)-1, 0, -1):\n",
        "        if r[fin_positive] > 0:\n",
        "            break\n",
        "    features[\"observations\"].append(o[:fin_positive+1])\n",
        "    features[\"actions\"].append(a[:fin_positive+1])\n",
        "    features[\"dones\"].append(d[:fin_positive+1])\n",
        "\n",
        "    r = r[:fin_positive+1]\n",
        "    features[\"rewards\"].append(r)    \n",
        "    rtg = np.zeros_like(r)\n",
        "    rtg[-1] = r[-1]\n",
        "    for i in reversed(range(len(r) - 1)):\n",
        "        rtg[i] = r[i] + RTG_GAMMA * rtg[i + 1]\n",
        "    features[\"rtg\"].append(rtg)\n",
        "    \n",
        "    print(f\"\\nTotal reward: {total_reward} steps: {len(o)}\")\n",
        "env.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [],
      "source": [
        "import utils.storage as storage\n",
        "dataset_file = f'car_racing_{NUM_EPISODES}_{MAX_EPISODE_STEPS}.hdf5'\n",
        "storage.save_to_hdf5(features, f\"./{dataset_file}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "#gcloud auth login\n",
        "!gcloud config set project clever-393810\n",
        "#!gcloud auth list\n",
        "#!gcloud storage buckets list\n",
        "!gsutil cp ./{dataset_file} gs://yakiv-dt-public/datasets/{dataset_file}"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "V100",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}