min_decision_transformer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "min_decision_transformer.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "Euib_RmfOiT3",
        "aBD3fRknjEj6",
        "9TpGEYTblzQc",
        "wNJM0LG1iziA",
        "gLHjV3q28LNr",
        "pewE01Ca4BG0",
        "QXXrs_PjAHrN",
        "wxcJqnb1Him4"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Euib_RmfOiT3"
      },
      "source": [
        "## install mujoco-py and D4RL\n",
        "\n",
        "* **Restart Runtime** after running this block to complete D4RL setup\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DAGCHznQs2bI"
      },
      "source": [
        "\n",
        "###### libs for install ######\n",
        "\n",
        "\n",
        "!sudo apt-get update\n",
        "!sudo apt-get install gcc\n",
        "\n",
        "!sudo apt-get build-dep mesa\n",
        "!sudo apt-get install llvm-dev\n",
        "!sudo apt-get install freeglut3 freeglut3-dev\n",
        "\n",
        "!sudo apt-get install python3-dev\n",
        "\n",
        "!sudo apt-get install build-essential\n",
        "\n",
        "!sudo apt install curl git libgl1-mesa-dev libgl1-mesa-glx libglew-dev \\\n",
        "        libosmesa6-dev software-properties-common net-tools unzip vim \\\n",
        "        virtualenv wget xpra xserver-xorg-dev libglfw3-dev patchelf\n",
        "\n",
        "#!sudo apt-get install -y libglew-dev\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "23eqLoV_orip"
      },
      "source": [
        "\n",
        "###### mujoco setup ######\n",
        "\n",
        "\n",
        "#!wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz\n",
        "\n",
        "!wget https://roboti.us/download/mujoco200_linux.zip\n",
        "\n",
        "!wget https://roboti.us/file/mjkey.txt\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WcXVniz_p4RN"
      },
      "source": [
        "\n",
        "!mkdir /root/.mujoco\n",
        "\n",
        "### mujoco 210\n",
        "#!tar -xf mujoco210-linux-x86_64.tar.gz -C /.mujoco/\n",
        "#!ls -alh /.mujoco/mujoco210\n",
        "\n",
        "### mujoco 200\n",
        "!unzip mujoco200_linux.zip -d /root/.mujoco/\n",
        "!cp -r /root/.mujoco/mujoco200_linux /root/.mujoco/mujoco200\n",
        "\n",
        "!mv mjkey.txt /root/.mujoco/\n",
        "\n",
        "!cp -r /root/.mujoco/mujoco200/bin/* /usr/lib/\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-X3JOM3RTcPO"
      },
      "source": [
        "\n",
        "!ls -alh /root/.mujoco/\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YVQWcww_uZMo"
      },
      "source": [
        "\n",
        "%env LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.mujoco/mujoco200/bin\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AviuDDxpqhOs"
      },
      "source": [
        "\n",
        "###### mujoco-py setup ######\n",
        "\n",
        "!pip install mujoco_py==2.0.2.8\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "duUbqfKEordx"
      },
      "source": [
        "\n",
        "###### D4RL setup ######\n",
        "\n",
        "## !pip uninstall dm_control==0.0.364896371\n",
        "\n",
        "!git clone https://github.com/rail-berkeley/d4rl.git\n",
        "\n",
        "### edit dm_control version in d4rl setup.py\n",
        "!sed -i \"s;dm_control @ git+git://github.com/deepmind/dm_control@master#egg=dm_control;dm_control==0.0.364896371;g\" /content/d4rl/setup.py\n",
        "\n",
        "### edit mjrl install in d4rl setup.py to use github's new https protocol instead of git SSH\n",
        "!sed -i \"s;mjrl @ git+git://github.com/aravindr93/mjrl@master#egg=mjrl;mjrl @ git+https://github.com/aravindr93/mjrl@master#egg=mjrl;g\" /content/d4rl/setup.py\n",
        "\n",
        "!pip install -e d4rl/.\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jVrmCbNMAwQk"
      },
      "source": [
        "\n",
        "###### restart runtime ######\n",
        "\n",
        "exit()\n"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# check mujoco-py and D4RL installation\n",
        "\n",
        "* if check fails then **Restart Runtime** again\n",
        "* if check still fails then Factory reset runtime and install again\n",
        "* After installing, first import will be slow as the lib will be built again\n"
      ],
      "metadata": {
        "id": "aBD3fRknjEj6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# set mujoco env path if not already set\n",
        "%env LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.mujoco/mujoco200/bin\n",
        "\n",
        "import gym\n",
        "import d4rl # Import required to register environments\n",
        "\n",
        "\n",
        "env = gym.make('Walker2d-v3')\n",
        "env.reset()\n",
        "env.step(env.action_space.sample())\n",
        "env.close()\n",
        "print(\"mujoco-py check passed\")\n",
        "\n",
        "env = gym.make('walker2d-medium-v2')\n",
        "env.reset()\n",
        "env.step(env.action_space.sample())\n",
        "env.close()\n",
        "print(\"d4rl check passed\")\n"
      ],
      "metadata": {
        "id": "3uycTGiqjKYK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# download D4RL data\n",
        "\n",
        "*   skip this block if data is already downloaded\n",
        "\n"
      ],
      "metadata": {
        "id": "9TpGEYTblzQc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# set mujoco env path if not already set\n",
        "%env LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.mujoco/mujoco200/bin\n",
        "\n",
        "import os\n",
        "import gym\n",
        "import numpy as np\n",
        "\n",
        "import collections\n",
        "import pickle\n",
        "\n",
        "import d4rl\n",
        "\n",
        "datasets = []\n",
        "\n",
        "data_dir = \"./data\"\n",
        "\n",
        "print(data_dir)\n",
        "\n",
        "if not os.path.exists(data_dir):\n",
        "    os.makedirs(data_dir)\n",
        "\n",
        "for env_name in ['walker2d', 'halfcheetah', 'hopper']:\n",
        "    for dataset_type in ['medium', 'medium-expert', 'medium-replay']:\n",
        "\t\t\n",
        "        name = f'{env_name}-{dataset_type}-v2'\n",
        "        pkl_file_path = os.path.join(data_dir, name)\n",
        "\n",
        "        print(\"processing: \", name)\n",
        "\n",
        "        env = gym.make(name)\n",
        "        dataset = env.get_dataset()\n",
        "\n",
        "        N = dataset['rewards'].shape[0]\n",
        "        data_ = collections.defaultdict(list)\n",
        "\n",
        "        use_timeouts = False\n",
        "        if 'timeouts' in dataset:\n",
        "            use_timeouts = True\n",
        "\n",
        "        episode_step = 0\n",
        "        paths = []\n",
        "        for i in range(N):\n",
        "            done_bool = bool(dataset['terminals'][i])\n",
        "            if use_timeouts:\n",
        "                final_timestep = dataset['timeouts'][i]\n",
        "            else:\n",
        "                final_timestep = (episode_step == 1000-1)\n",
        "            for k in ['observations', 'next_observations', 'actions', 'rewards', 'terminals']:\n",
        "                data_[k].append(dataset[k][i])\n",
        "            if done_bool or final_timestep:\n",
        "                episode_step = 0\n",
        "                episode_data = {}\n",
        "                for k in data_:\n",
        "                    episode_data[k] = np.array(data_[k])\n",
        "                paths.append(episode_data)\n",
        "                data_ = collections.defaultdict(list)\n",
        "            episode_step += 1\n",
        "\n",
        "        returns = np.array([np.sum(p['rewards']) for p in paths])\n",
        "        num_samples = np.sum([p['rewards'].shape[0] for p in paths])\n",
        "        print(f'Number of samples collected: {num_samples}')\n",
        "        print(f'Trajectory returns: mean = {np.mean(returns)}, std = {np.std(returns)}, max = {np.max(returns)}, min = {np.min(returns)}')\n",
        "\n",
        "        with open(f'{pkl_file_path}.pkl', 'wb') as f:\n",
        "            pickle.dump(paths, f)\n",
        "\n"
      ],
      "metadata": {
        "id": "V31ELEKOih7D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "kHGzjOmbamJz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# import libs\n"
      ],
      "metadata": {
        "id": "W2Nk5Gp7hUGA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# set mujoco env path if not already set\n",
        "%env LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.mujoco/mujoco200/bin\n",
        "\n",
        "import os\n",
        "import sys\n",
        "import random\n",
        "import csv\n",
        "from datetime import datetime\n",
        "import pickle\n",
        "import collections\n",
        "import math\n",
        "\n",
        "import numpy as np\n",
        "import gym\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "\n"
      ],
      "metadata": {
        "id": "q4xiijmBixUm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# training parameters"
      ],
      "metadata": {
        "id": "QQcLNRgD6SaW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "dataset = \"medium\"       # medium / medium-replay / medium-expert\n",
        "rtg_scale = 1000                # scale to normalize returns to go\n",
        "\n",
        "# use v3 env for evaluation because\n",
        "# DT paper evaluates results on v3 envs\n",
        "\n",
        "env_name = 'Walker2d-v3'\n",
        "rtg_target = 5000\n",
        "env_d4rl_name = f'walker2d-{dataset}-v2'\n",
        "\n",
        "# env_name = 'HalfCheetah-v3'\n",
        "# rtg_target = 6000\n",
        "# env_d4rl_name = f'halfcheetah-{dataset}-v2'\n",
        "\n",
        "# env_name = 'Hopper-v3'\n",
        "# rtg_target = 3600\n",
        "# env_d4rl_name = f'hopper-{dataset}-v2'\n",
        "\n",
        "\n",
        "max_eval_ep_len = 1000      # max len of one evaluation episode\n",
        "num_eval_ep = 10            # num of evaluation episodes per iteration\n",
        "\n",
        "batch_size = 64             # training batch size\n",
        "lr = 1e-4                   # learning rate\n",
        "wt_decay = 1e-4             # weight decay\n",
        "warmup_steps = 10000        # warmup steps for lr scheduler\n",
        "\n",
        "# total updates = max_train_iters x num_updates_per_iter\n",
        "max_train_iters = 200\n",
        "num_updates_per_iter = 100\n",
        "\n",
        "context_len = 20        # K in decision transformer\n",
        "n_blocks = 3            # num of transformer blocks\n",
        "embed_dim = 128         # embedding (hidden) dim of transformer\n",
        "n_heads = 1             # num of transformer heads\n",
        "dropout_p = 0.1         # dropout probability\n",
        "\n",
        "\n",
        "\n",
        "# load data from this file\n",
        "dataset_path = f'data/{env_d4rl_name}.pkl'\n",
        "\n",
        "# saves model and csv in this directory\n",
        "log_dir = \"./dt_runs/\"\n",
        "\n",
        "\n",
        "if not os.path.exists(log_dir):\n",
        "    os.makedirs(log_dir)\n",
        "\n",
        "\n",
        "# training and evaluation device\n",
        "device_name = 'cuda'\n",
        "device = torch.device(device_name)\n",
        "print(\"device set to: \", device)\n",
        "\n"
      ],
      "metadata": {
        "id": "WdtDsvit6m_e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# decision transformer model"
      ],
      "metadata": {
        "id": "wNJM0LG1iziA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\"\"\"\n",
        "this extremely minimal GPT model is based on:\n",
        "Misha Laskin's tweet: \n",
        "https://twitter.com/MishaLaskin/status/1481767788775628801?cxt=HHwWgoCzmYD9pZApAAAA\n",
        "\n",
        "and its corresponding notebook:\n",
        "https://colab.research.google.com/drive/1NUBqyboDcGte5qAJKOl8gaJC28V_73Iv?usp=sharing\n",
        "\n",
        "the above colab has a bug while applying masked_fill which is fixed in the\n",
        "following code\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "class MaskedCausalAttention(nn.Module):\n",
        "    def __init__(self, h_dim, max_T, n_heads, drop_p):\n",
        "        super().__init__()\n",
        "\n",
        "        self.n_heads = n_heads\n",
        "        self.max_T = max_T\n",
        "\n",
        "        self.q_net = nn.Linear(h_dim, h_dim)\n",
        "        self.k_net = nn.Linear(h_dim, h_dim)\n",
        "        self.v_net = nn.Linear(h_dim, h_dim)\n",
        "\n",
        "        self.proj_net = nn.Linear(h_dim, h_dim)\n",
        "\n",
        "        self.att_drop = nn.Dropout(drop_p)\n",
        "        self.proj_drop = nn.Dropout(drop_p)\n",
        "\n",
        "        ones = torch.ones((max_T, max_T))\n",
        "        mask = torch.tril(ones).view(1, 1, max_T, max_T)\n",
        "\n",
        "        # register buffer makes sure mask does not get updated\n",
        "        # during backpropagation\n",
        "        self.register_buffer('mask',mask)\n",
        "\n",
        "    def forward(self, x):\n",
        "        B, T, C = x.shape # batch size, seq length, h_dim * n_heads\n",
        "\n",
        "        N, D = self.n_heads, C // self.n_heads # N = num heads, D = attention dim\n",
        "\n",
        "        # rearrange q, k, v as (B, N, T, D)\n",
        "        q = self.q_net(x).view(B, T, N, D).transpose(1,2)\n",
        "        k = self.k_net(x).view(B, T, N, D).transpose(1,2)\n",
        "        v = self.v_net(x).view(B, T, N, D).transpose(1,2)\n",
        "\n",
        "        # weights (B, N, T, T)\n",
        "        weights = q @ k.transpose(2,3) / math.sqrt(D)\n",
        "        # causal mask applied to weights\n",
        "        weights = weights.masked_fill(self.mask[...,:T,:T] == 0, float('-inf'))\n",
        "        # normalize weights, all -inf -> 0 after softmax\n",
        "        normalized_weights = F.softmax(weights, dim=-1)\n",
        "\n",
        "        # attention (B, N, T, D)\n",
        "        attention = self.att_drop(normalized_weights @ v)\n",
        "\n",
        "        # gather heads and project (B, N, T, D) -> (B, T, N*D)\n",
        "        attention = attention.transpose(1, 2).contiguous().view(B,T,N*D)\n",
        "\n",
        "        out = self.proj_drop(self.proj_net(attention))\n",
        "        return out\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, h_dim, max_T, n_heads, drop_p):\n",
        "        super().__init__()\n",
        "        self.attention = MaskedCausalAttention(h_dim, max_T, n_heads, drop_p)\n",
        "        self.mlp = nn.Sequential(\n",
        "                nn.Linear(h_dim, 4*h_dim),\n",
        "                nn.GELU(),\n",
        "                nn.Linear(4*h_dim, h_dim),\n",
        "                nn.Dropout(drop_p),\n",
        "            )\n",
        "        self.ln1 = nn.LayerNorm(h_dim)\n",
        "        self.ln2 = nn.LayerNorm(h_dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Attention -> LayerNorm -> MLP -> LayerNorm\n",
        "        x = x + self.attention(x) # residual\n",
        "        x = self.ln1(x)\n",
        "        x = x + self.mlp(x) # residual\n",
        "        x = self.ln2(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class DecisionTransformer(nn.Module):\n",
        "    def __init__(self, state_dim, act_dim, n_blocks, h_dim, context_len, \n",
        "                 n_heads, drop_p, max_timestep=4096):\n",
        "        super().__init__()\n",
        "\n",
        "        self.state_dim = state_dim\n",
        "        self.act_dim = act_dim\n",
        "        self.h_dim = h_dim\n",
        "\n",
        "        ### transformer blocks\n",
        "        input_seq_len = 3 * context_len\n",
        "        blocks = [Block(h_dim, input_seq_len, n_heads, drop_p) for _ in range(n_blocks)]\n",
        "        self.transformer = nn.Sequential(*blocks)\n",
        "\n",
        "        ### projection heads (project to embedding)\n",
        "        self.embed_ln = nn.LayerNorm(h_dim)\n",
        "        self.embed_timestep = nn.Embedding(max_timestep, h_dim)\n",
        "        self.embed_rtg = torch.nn.Linear(1, h_dim)\n",
        "        self.embed_state = torch.nn.Linear(state_dim, h_dim)\n",
        "        \n",
        "        # # discrete actions\n",
        "        # self.embed_action = torch.nn.Embedding(act_dim, h_dim)\n",
        "        # use_action_tanh = False # False for discrete actions\n",
        "\n",
        "        # continuous actions\n",
        "        self.embed_action = torch.nn.Linear(act_dim, h_dim)\n",
        "        use_action_tanh = True # True for continuous actions\n",
        "        \n",
        "        ### prediction heads\n",
        "        self.predict_rtg = torch.nn.Linear(h_dim, 1)\n",
        "        self.predict_state = torch.nn.Linear(h_dim, state_dim)\n",
        "        self.predict_action = nn.Sequential(\n",
        "            *([nn.Linear(h_dim, act_dim)] + ([nn.Tanh()] if use_action_tanh else []))\n",
        "        )\n",
        "\n",
        "\n",
        "    def forward(self, timesteps, states, actions, returns_to_go):\n",
        "\n",
        "        B, T, _ = states.shape\n",
        "\n",
        "        time_embeddings = self.embed_timestep(timesteps)\n",
        "\n",
        "        # time embeddings are treated similar to positional embeddings\n",
        "        state_embeddings = self.embed_state(states) + time_embeddings\n",
        "        action_embeddings = self.embed_action(actions) + time_embeddings\n",
        "        returns_embeddings = self.embed_rtg(returns_to_go) + time_embeddings\n",
        "\n",
        "        # stack rtg, states and actions and reshape sequence as\n",
        "        # (r1, s1, a1, r2, s2, a2 ...)\n",
        "        h = torch.stack(\n",
        "            (returns_embeddings, state_embeddings, action_embeddings), dim=1\n",
        "        ).permute(0, 2, 1, 3).reshape(B, 3 * T, self.h_dim)\n",
        "\n",
        "        h = self.embed_ln(h)\n",
        "        \n",
        "        # transformer and prediction\n",
        "        h = self.transformer(h)\n",
        "\n",
        "        # get h reshaped such that its size = (B x 3 x T x h_dim) and\n",
        "        # h[:, 0, t] is conditioned on r_0, s_0, a_0 ... r_t\n",
        "        # h[:, 1, t] is conditioned on r_0, s_0, a_0 ... r_t, s_t\n",
        "        # h[:, 2, t] is conditioned on r_0, s_0, a_0 ... r_t, s_t, a_t\n",
        "        h = h.reshape(B, T, 3, self.h_dim).permute(0, 2, 1, 3)\n",
        "\n",
        "        # get predictions\n",
        "        return_preds = self.predict_rtg(h[:,2])     # predict next rtg given r, s, a\n",
        "        state_preds = self.predict_state(h[:,2])    # predict next state given r, s, a\n",
        "        action_preds = self.predict_action(h[:,1])  # predict action given r, s\n",
        "    \n",
        "        return state_preds, action_preds, return_preds\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "MHMl_Y1SicXb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# infos"
      ],
      "metadata": {
        "id": "gLHjV3q28LNr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "## from infos.py from official d4rl github repo\n",
        "\n",
        "REF_MAX_SCORE = {\n",
        "    'halfcheetah' : 12135.0,\n",
        "    'walker2d' : 4592.3,\n",
        "    'hopper' : 3234.3,\n",
        "}\n",
        "\n",
        "REF_MIN_SCORE = {\n",
        "    'halfcheetah' : -280.178953,\n",
        "    'walker2d' : 1.629008,\n",
        "    'hopper' : -20.272305,\n",
        "}\n",
        "\n",
        "\n",
        "## calculated from d4rl datasets\n",
        "\n",
        "D4RL_DATASET_STATS = {\n",
        "        'halfcheetah-medium-v2': {\n",
        "                'state_mean':[-0.06845773756504059, 0.016414547339081764, -0.18354906141757965, \n",
        "                              -0.2762460708618164, -0.34061527252197266, -0.09339715540409088, \n",
        "                              -0.21321271359920502, -0.0877423882484436, 5.173007488250732, \n",
        "                              -0.04275195300579071, -0.036108363419771194, 0.14053793251514435, \n",
        "                              0.060498327016830444, 0.09550975263118744, 0.06739100068807602, \n",
        "                              0.005627387668937445, 0.013382787816226482\n",
        "                ],\n",
        "                'state_std':[0.07472999393939972, 0.3023499846458435, 0.30207309126853943, \n",
        "                             0.34417077898979187, 0.17619241774082184, 0.507205605506897, \n",
        "                             0.2567007839679718, 0.3294812738895416, 1.2574149370193481, \n",
        "                             0.7600541710853577, 1.9800915718078613, 6.565362453460693, \n",
        "                             7.466367721557617, 4.472222805023193, 10.566964149475098, \n",
        "                             5.671932697296143, 7.4982590675354  \n",
        "                ]\n",
        "            },\n",
        "        'halfcheetah-medium-replay-v2': {\n",
        "                'state_mean':[-0.12880703806877136, 0.3738119602203369, -0.14995987713336945, \n",
        "                              -0.23479078710079193, -0.2841278612613678, -0.13096535205841064, \n",
        "                              -0.20157982409000397, -0.06517726927995682, 3.4768247604370117, \n",
        "                              -0.02785065770149231, -0.015035249292850494, 0.07697279006242752, \n",
        "                              0.01266712136566639, 0.027325302362442017, 0.02316424623131752, \n",
        "                              0.010438721626996994, -0.015839405357837677\n",
        "                ],\n",
        "                'state_std':[0.17019015550613403, 1.284424901008606, 0.33442774415016174, \n",
        "                             0.3672759234905243, 0.26092398166656494, 0.4784106910228729, \n",
        "                             0.3181420564651489, 0.33552637696266174, 2.0931615829467773, \n",
        "                             0.8037433624267578, 1.9044333696365356, 6.573209762573242, \n",
        "                             7.572863578796387, 5.069749355316162, 9.10555362701416, \n",
        "                             6.085654258728027, 7.25300407409668\n",
        "                ]\n",
        "            },\n",
        "        'halfcheetah-medium-expert-v2': {\n",
        "                'state_mean':[-0.05667462572455406, 0.024369969964027405, -0.061670560389757156, \n",
        "                              -0.22351515293121338, -0.2675151228904724, -0.07545716315507889, \n",
        "                              -0.05809682980179787, -0.027675075456500053, 8.110626220703125, \n",
        "                              -0.06136331334710121, -0.17986927926540375, 0.25175222754478455, \n",
        "                              0.24186332523822784, 0.2519369423389435, 0.5879552960395813, \n",
        "                              -0.24090635776519775, -0.030184272676706314\n",
        "                ],\n",
        "                'state_std':[0.06103534251451492, 0.36054104566574097, 0.45544400811195374, \n",
        "                             0.38476887345314026, 0.2218363732099533, 0.5667523741722107, \n",
        "                             0.3196682929992676, 0.2852923572063446, 3.443821907043457, \n",
        "                             0.6728139519691467, 1.8616976737976074, 9.575807571411133, \n",
        "                             10.029894828796387, 5.903450012207031, 12.128185272216797, \n",
        "                             6.4811787605285645, 6.378620147705078\n",
        "                ]\n",
        "            },\n",
        "        'walker2d-medium-v2': {\n",
        "                'state_mean':[1.218966007232666, 0.14163373410701752, -0.03704913705587387, \n",
        "                              -0.13814310729503632, 0.5138224363327026, -0.04719110205769539, \n",
        "                              -0.47288352251052856, 0.042254164814949036, 2.3948874473571777, \n",
        "                              -0.03143199160695076, 0.04466355964541435, -0.023907244205474854, \n",
        "                              -0.1013401448726654, 0.09090937674045563, -0.004192637279629707, \n",
        "                              -0.12120571732521057, -0.5497063994407654\n",
        "                ],\n",
        "                'state_std':[0.12311358004808426, 0.3241879940032959, 0.11456084251403809, \n",
        "                             0.2623065710067749, 0.5640279054641724, 0.2271878570318222, \n",
        "                             0.3837319612503052, 0.7373676896095276, 1.2387926578521729, \n",
        "                             0.798020601272583, 1.5664079189300537, 1.8092705011367798, \n",
        "                             3.025604248046875, 4.062486171722412, 1.4586567878723145, \n",
        "                             3.7445690631866455, 5.5851287841796875\n",
        "                ]\n",
        "            },\n",
        "        'walker2d-medium-replay-v2': {\n",
        "                'state_mean':[1.209364652633667, 0.13264022767543793, -0.14371201395988464, \n",
        "                              -0.2046516090631485, 0.5577612519264221, -0.03231537342071533, \n",
        "                              -0.2784661054611206, 0.19130706787109375, 1.4701707363128662, \n",
        "                              -0.12504704296588898, 0.0564953051507473, -0.09991033375263214, \n",
        "                              -0.340340256690979, 0.03546293452382088, -0.08934258669614792, \n",
        "                              -0.2992438077926636, -0.5984178185462952   \n",
        "                ],\n",
        "                'state_std':[0.11929835379123688, 0.3562574088573456, 0.25852200388908386, \n",
        "                             0.42075422406196594, 0.5202291011810303, 0.15685082972049713, \n",
        "                             0.36770978569984436, 0.7161387801170349, 1.3763766288757324, \n",
        "                             0.8632221817970276, 2.6364643573760986, 3.0134117603302, \n",
        "                             3.720684051513672, 4.867283821105957, 2.6681625843048096, \n",
        "                             3.845186948776245, 5.4768385887146\n",
        "                ]\n",
        "            },\n",
        "        'walker2d-medium-expert-v2': {\n",
        "                'state_mean':[1.2294334173202515, 0.16869689524173737, -0.07089081406593323, \n",
        "                              -0.16197483241558075, 0.37101927399635315, -0.012209027074277401, \n",
        "                              -0.42461398243904114, 0.18986578285694122, 3.162475109100342, \n",
        "                              -0.018092676997184753, 0.03496946766972542, -0.013921679928898811, \n",
        "                              -0.05937029421329498, -0.19549426436424255, -0.0019200450042262673, \n",
        "                              -0.062483321875333786, -0.27366524934768677\n",
        "                ],\n",
        "                'state_std':[0.09932824969291687, 0.25981399416923523, 0.15062759816646576, \n",
        "                             0.24249176681041718, 0.6758718490600586, 0.1650741547346115, \n",
        "                             0.38140663504600525, 0.6962361335754395, 1.3501490354537964, \n",
        "                             0.7641991376876831, 1.534574270248413, 2.1785972118377686, \n",
        "                             3.276582717895508, 4.766193866729736, 1.1716983318328857, \n",
        "                             4.039782524108887, 5.891613960266113       \n",
        "                ]\n",
        "            },\n",
        "        'hopper-medium-v2': {\n",
        "                'state_mean':[1.311279058456421, -0.08469521254301071, -0.5382719039916992, \n",
        "                              -0.07201576232910156, 0.04932365566492081, 2.1066856384277344, \n",
        "                              -0.15017354488372803, 0.008783451281487942, -0.2848185896873474, \n",
        "                              -0.18540096282958984, -0.28461286425590515\n",
        "                ],\n",
        "                'state_std':[0.17790751159191132, 0.05444620922207832, 0.21297138929367065, \n",
        "                             0.14530418813228607, 0.6124444007873535, 0.8517446517944336, \n",
        "                             1.4515252113342285, 0.6751695871353149, 1.5362390279769897, \n",
        "                             1.616074562072754, 5.607253551483154 \n",
        "                ]\n",
        "            },\n",
        "        'hopper-medium-replay-v2': {\n",
        "                'state_mean':[1.2305138111114502, -0.04371410980820656, -0.44542956352233887, \n",
        "                              -0.09370097517967224, 0.09094487875699997, 1.3694725036621094, \n",
        "                              -0.19992674887180328, -0.022861352190375328, -0.5287045240402222, \n",
        "                              -0.14465883374214172, -0.19652697443962097      \n",
        "                ],\n",
        "                'state_std':[0.1756512075662613, 0.0636928603053093, 0.3438323438167572, \n",
        "                             0.19566889107227325, 0.5547984838485718, 1.051029920578003, \n",
        "                             1.158307671546936, 0.7963128685951233, 1.4802359342575073, \n",
        "                             1.6540331840515137, 5.108601093292236\n",
        "                ]\n",
        "            },\n",
        "        'hopper-medium-expert-v2': {\n",
        "                'state_mean':[1.3293815851211548, -0.09836531430482864, -0.5444297790527344, \n",
        "                              -0.10201650857925415, 0.02277466468513012, 2.3577215671539307, \n",
        "                              -0.06349576264619827, -0.00374026270583272, -0.1766270101070404, \n",
        "                              -0.11862941086292267, -0.12097819894552231\n",
        "                ],\n",
        "                'state_std':[0.17012375593185425, 0.05159067362546921, 0.18141433596611023, \n",
        "                             0.16430604457855225, 0.6023368239402771, 0.7737284898757935, \n",
        "                             1.4986555576324463, 0.7483318448066711, 1.7953159809112549, \n",
        "                             2.0530025959014893, 5.725032806396484\n",
        "                ]\n",
        "            },\n",
        "    }\n",
        "\n"
      ],
      "metadata": {
        "id": "btnq_IL_j4PO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# utils"
      ],
      "metadata": {
        "id": "pewE01Ca4BG0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "def discount_cumsum(x, gamma):\n",
        "    disc_cumsum = np.zeros_like(x)\n",
        "    disc_cumsum[-1] = x[-1]\n",
        "    for t in reversed(range(x.shape[0]-1)):\n",
        "        disc_cumsum[t] = x[t] + gamma * disc_cumsum[t+1]\n",
        "    return disc_cumsum\n",
        "\n",
        "\n",
        "def get_d4rl_dataset_stats(env_d4rl_name):\n",
        "    return D4RL_DATASET_STATS[env_d4rl_name]\n",
        "\n",
        "\n",
        "def get_d4rl_normalized_score(score, env_name):\n",
        "    env_key = env_name.split('-')[0].lower()\n",
        "    assert env_key in REF_MAX_SCORE, f'no reference score for {env_key} env to calculate d4rl score'\n",
        "    return (score - REF_MIN_SCORE[env_key]) / (REF_MAX_SCORE[env_key] - REF_MIN_SCORE[env_key])\n",
        "    \n",
        "    \n",
        "def evaluate_on_env(model, device, context_len, env, rtg_target, rtg_scale,\n",
        "                    num_eval_ep=10, max_test_ep_len=1000,\n",
        "                    state_mean=None, state_std=None, render=False):\n",
        "\n",
        "    eval_batch_size = 1  # required for forward pass\n",
        "\n",
        "    results = {}\n",
        "    total_reward = 0\n",
        "    total_timesteps = 0\n",
        "\n",
        "    state_dim = env.observation_space.shape[0]\n",
        "    act_dim = env.action_space.shape[0]\n",
        "\n",
        "    if state_mean is None:\n",
        "        state_mean = torch.zeros((state_dim,)).to(device)\n",
        "    else:\n",
        "        state_mean = torch.from_numpy(state_mean).to(device)\n",
        "        \n",
        "    if state_std is None:\n",
        "        state_std = torch.ones((state_dim,)).to(device)\n",
        "    else:\n",
        "        state_std = torch.from_numpy(state_std).to(device)\n",
        "\n",
        "    # same as timesteps used for training the transformer\n",
        "    # also, crashes if device is passed to arange()\n",
        "    timesteps = torch.arange(start=0, end=max_test_ep_len, step=1)\n",
        "    timesteps = timesteps.repeat(eval_batch_size, 1).to(device)\n",
        "\n",
        "    model.eval()\n",
        "\n",
        "    with torch.no_grad():\n",
        "\n",
        "        for _ in range(num_eval_ep):\n",
        "\n",
        "            # zeros place holders\n",
        "            actions = torch.zeros((eval_batch_size, max_test_ep_len, act_dim),\n",
        "                                dtype=torch.float32, device=device)\n",
        "\n",
        "            states = torch.zeros((eval_batch_size, max_test_ep_len, state_dim),\n",
        "                                dtype=torch.float32, device=device)\n",
        "            \n",
        "            rewards_to_go = torch.zeros((eval_batch_size, max_test_ep_len, 1),\n",
        "                                dtype=torch.float32, device=device)\n",
        "            \n",
        "            # init episode\n",
        "            running_state = env.reset()\n",
        "            running_reward = 0\n",
        "            running_rtg = rtg_target / rtg_scale\n",
        "\n",
        "            for t in range(max_test_ep_len):\n",
        "\n",
        "                total_timesteps += 1\n",
        "\n",
        "                # add state in placeholder and normalize\n",
        "                states[0, t] = torch.from_numpy(running_state).to(device)\n",
        "                states[0, t] = (states[0, t] - state_mean) / state_std\n",
        "\n",
        "                # calcualate running rtg and add in placeholder\n",
        "                running_rtg = running_rtg - (running_reward / rtg_scale)\n",
        "                rewards_to_go[0, t] = running_rtg\n",
        "\n",
        "                if t < context_len:\n",
        "                    _, act_preds, _ = model.forward(timesteps[:,:context_len],\n",
        "                                                states[:,:context_len],\n",
        "                                                actions[:,:context_len],\n",
        "                                                rewards_to_go[:,:context_len])\n",
        "                    act = act_preds[0, t].detach()\n",
        "                else:\n",
        "                    _, act_preds, _ = model.forward(timesteps[:,t-context_len+1:t+1],\n",
        "                                                states[:,t-context_len+1:t+1],\n",
        "                                                actions[:,t-context_len+1:t+1],\n",
        "                                                rewards_to_go[:,t-context_len+1:t+1])\n",
        "                    act = act_preds[0, -1].detach()\n",
        "\n",
        "\n",
        "                running_state, running_reward, done, _ = env.step(act.cpu().numpy())\n",
        "\n",
        "                # add action in placeholder\n",
        "                actions[0, t] = act\n",
        "\n",
        "                total_reward += running_reward\n",
        "\n",
        "                if render:\n",
        "                    env.render()\n",
        "                if done:\n",
        "                    break\n",
        "\n",
        "    results['eval/avg_reward'] = total_reward / num_eval_ep\n",
        "    results['eval/avg_ep_len'] = total_timesteps / num_eval_ep\n",
        "    \n",
        "    return results\n",
        "\n"
      ],
      "metadata": {
        "id": "EaaymCHPlynF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "KDk9X9jJ8iAZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# dataset"
      ],
      "metadata": {
        "id": "QXXrs_PjAHrN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## check data\n",
        "\n",
        "# load dataset\n",
        "with open(dataset_path, 'rb') as f:\n",
        "    trajectories = pickle.load(f)\n",
        "\n",
        "min_len = 10**4\n",
        "states = []\n",
        "for traj in trajectories:\n",
        "    min_len = min(min_len, traj['observations'].shape[0])\n",
        "    states.append(traj['observations'])\n",
        "\n",
        "# used for input normalization\n",
        "states = np.concatenate(states, axis=0)\n",
        "state_mean, state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6\n",
        "\n",
        "print(dataset_path)\n",
        "print(\"num of trajectories in dataset: \", len(trajectories))\n",
        "print(\"minimum trajectory length in dataset: \", min_len)\n",
        "print(\"state mean: \", state_mean.tolist())\n",
        "print(\"state std: \", state_std.tolist())\n",
        "\n",
        "\n",
        "## check if info is correct\n",
        "print(\"is state mean info correct: \", state_mean.tolist() == D4RL_DATASET_STATS[env_d4rl_name]['state_mean'])\n",
        "print(\"is state std info correct: \", state_std.tolist() == D4RL_DATASET_STATS[env_d4rl_name]['state_std'])\n",
        "\n",
        "\n",
        "assert state_mean.tolist() == D4RL_DATASET_STATS[env_d4rl_name]['state_mean']\n",
        "assert state_std.tolist() == D4RL_DATASET_STATS[env_d4rl_name]['state_std']\n"
      ],
      "metadata": {
        "id": "n1Vb5rY_iiME"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "class D4RLTrajectoryDataset(Dataset):\n",
        "    def __init__(self, dataset_path, context_len, rtg_scale):\n",
        "\n",
        "        self.context_len = context_len        \n",
        "\n",
        "        # load dataset\n",
        "        with open(dataset_path, 'rb') as f:\n",
        "            self.trajectories = pickle.load(f)\n",
        "        \n",
        "        # calculate min len of traj, state mean and variance\n",
        "        # and returns_to_go for all traj\n",
        "        min_len = 10**6\n",
        "        states = []\n",
        "        for traj in self.trajectories:\n",
        "            traj_len = traj['observations'].shape[0]\n",
        "            min_len = min(min_len, traj_len)\n",
        "            states.append(traj['observations'])\n",
        "            # calculate returns to go and rescale them\n",
        "            traj['returns_to_go'] = discount_cumsum(traj['rewards'], 1.0) / rtg_scale\n",
        "\n",
        "        # used for input normalization\n",
        "        states = np.concatenate(states, axis=0)\n",
        "        self.state_mean, self.state_std = np.mean(states, axis=0), np.std(states, axis=0) + 1e-6\n",
        "\n",
        "        # normalize states\n",
        "        for traj in self.trajectories:\n",
        "            traj['observations'] = (traj['observations'] - self.state_mean) / self.state_std\n",
        "\n",
        "\n",
        "    def get_state_stats(self):\n",
        "        return self.state_mean, self.state_std\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.trajectories)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        traj = self.trajectories[idx]\n",
        "        traj_len = traj['observations'].shape[0]\n",
        "\n",
        "        if traj_len >= self.context_len:\n",
        "            # sample random index to slice trajectory\n",
        "            si = random.randint(0, traj_len - self.context_len)\n",
        "\n",
        "            states = torch.from_numpy(traj['observations'][si : si + self.context_len])\n",
        "            actions = torch.from_numpy(traj['actions'][si : si + self.context_len])\n",
        "            returns_to_go = torch.from_numpy(traj['returns_to_go'][si : si + self.context_len])\n",
        "            timesteps = torch.arange(start=si, end=si+self.context_len, step=1)\n",
        "\n",
        "            # all ones since no padding\n",
        "            traj_mask = torch.ones(self.context_len, dtype=torch.long)\n",
        "\n",
        "        else:\n",
        "            padding_len = self.context_len - traj_len\n",
        "\n",
        "            # padding with zeros\n",
        "            states = torch.from_numpy(traj['observations'])\n",
        "            states = torch.cat([states,\n",
        "                                torch.zeros(([padding_len] + list(states.shape[1:])),\n",
        "                                dtype=states.dtype)], \n",
        "                               dim=0)\n",
        "            \n",
        "            actions = torch.from_numpy(traj['actions'])\n",
        "            actions = torch.cat([actions,\n",
        "                                torch.zeros(([padding_len] + list(actions.shape[1:])),\n",
        "                                dtype=actions.dtype)], \n",
        "                               dim=0)\n",
        "\n",
        "            returns_to_go = torch.from_numpy(traj['returns_to_go'])\n",
        "            returns_to_go = torch.cat([returns_to_go,\n",
        "                                torch.zeros(([padding_len] + list(returns_to_go.shape[1:])),\n",
        "                                dtype=returns_to_go.dtype)], \n",
        "                               dim=0)\n",
        "            \n",
        "            timesteps = torch.arange(start=0, end=self.context_len, step=1)\n",
        "\n",
        "            traj_mask = torch.cat([torch.ones(traj_len, dtype=torch.long), \n",
        "                                   torch.zeros(padding_len, dtype=torch.long)], \n",
        "                                  dim=0)\n",
        "            \n",
        "        return  timesteps, states, actions, returns_to_go, traj_mask\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "eo4zPTjjn0Qr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# train"
      ],
      "metadata": {
        "id": "p7AK6T9Picu-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "start_time = datetime.now().replace(microsecond=0)\n",
        "\n",
        "start_time_str = start_time.strftime(\"%y-%m-%d-%H-%M-%S\")\n",
        "\n",
        "prefix = \"dt_\" + env_d4rl_name\n",
        "\n",
        "save_model_name =  prefix + \"_model_\" + start_time_str + \".pt\"\n",
        "save_model_path = os.path.join(log_dir, save_model_name)\n",
        "save_best_model_path = save_model_path[:-3] + \"_best.pt\"\n",
        "\n",
        "log_csv_name = prefix + \"_log_\" + start_time_str + \".csv\"\n",
        "log_csv_path = os.path.join(log_dir, log_csv_name)\n",
        "\n",
        "\n",
        "csv_writer = csv.writer(open(log_csv_path, 'a', 1))\n",
        "csv_header = ([\"duration\", \"num_updates\", \"action_loss\", \n",
        "               \"eval_avg_reward\", \"eval_avg_ep_len\", \"eval_d4rl_score\"])\n",
        "\n",
        "csv_writer.writerow(csv_header)\n",
        "\n",
        "\n",
        "print(\"=\" * 60)\n",
        "print(\"start time: \" + start_time_str)\n",
        "print(\"=\" * 60)\n",
        "\n",
        "print(\"device set to: \" + str(device))\n",
        "print(\"dataset path: \" + dataset_path)\n",
        "print(\"model save path: \" + save_model_path)\n",
        "print(\"log csv save path: \" + log_csv_path)\n",
        "\n",
        "\n",
        "traj_dataset = D4RLTrajectoryDataset(dataset_path, context_len, rtg_scale)\n",
        "\n",
        "traj_data_loader = DataLoader(traj_dataset,\n",
        "\t\t\t\t\t\tbatch_size=batch_size,\n",
        "\t\t\t\t\t\tshuffle=True,\n",
        "\t\t\t\t\t\tpin_memory=True,\n",
        "\t\t\t\t\t\tdrop_last=True) \n",
        "\n",
        "data_iter = iter(traj_data_loader)\n",
        "\n",
        "## get state stats from dataset\n",
        "state_mean, state_std = traj_dataset.get_state_stats()\n",
        "\n",
        "env = gym.make(env_name)\n",
        "\n",
        "state_dim = env.observation_space.shape[0]\n",
        "act_dim = env.action_space.shape[0]\n",
        "\n",
        "model = DecisionTransformer(\n",
        "\t\t\tstate_dim=state_dim,\n",
        "\t\t\tact_dim=act_dim,\n",
        "\t\t\tn_blocks=n_blocks,\n",
        "\t\t\th_dim=embed_dim,\n",
        "\t\t\tcontext_len=context_len,\n",
        "\t\t\tn_heads=n_heads,\n",
        "\t\t\tdrop_p=dropout_p,\n",
        "\t\t).to(device)\n",
        "  \n",
        "optimizer = torch.optim.AdamW(\n",
        "\t\t\t\t\tmodel.parameters(), \n",
        "\t\t\t\t\tlr=lr, \n",
        "\t\t\t\t\tweight_decay=wt_decay\n",
        "\t\t\t\t)\n",
        "\n",
        "scheduler = torch.optim.lr_scheduler.LambdaLR(\n",
        "\t\toptimizer,\n",
        "\t\tlambda steps: min((steps+1)/warmup_steps, 1)\n",
        "\t)\n",
        "\n",
        "max_d4rl_score = -1.0\n",
        "total_updates = 0\n",
        "\n",
        "for i_train_iter in range(max_train_iters):\n",
        "\n",
        "\tlog_action_losses = []\t\n",
        "\tmodel.train()\n",
        " \n",
        "\tfor _ in range(num_updates_per_iter):\n",
        "\t\ttry:\n",
        "\t\t\ttimesteps, states, actions, returns_to_go, traj_mask = next(data_iter)\n",
        "\t\texcept StopIteration:\n",
        "\t\t\tdata_iter = iter(traj_data_loader)\n",
        "\t\t\ttimesteps, states, actions, returns_to_go, traj_mask = next(data_iter)\n",
        "\n",
        "\t\ttimesteps = timesteps.to(device)\t# B x T\n",
        "\t\tstates = states.to(device)\t\t\t# B x T x state_dim\n",
        "\t\tactions = actions.to(device)\t\t# B x T x act_dim\n",
        "\t\treturns_to_go = returns_to_go.to(device).unsqueeze(dim=-1) # B x T x 1\n",
        "\t\ttraj_mask = traj_mask.to(device)\t# B x T\n",
        "\n",
        "\t\taction_target = torch.clone(actions).detach().to(device)\n",
        "\t\n",
        "\t\tstate_preds, action_preds, return_preds = model.forward(\n",
        "\t\t\t\t\t\t\t\t\t\t\t\t\t\ttimesteps=timesteps,\n",
        "\t\t\t\t\t\t\t\t\t\t\t\t\t\tstates=states,\n",
        "\t\t\t\t\t\t\t\t\t\t\t\t\t\tactions=actions,\n",
        "\t\t\t\t\t\t\t\t\t\t\t\t\t\treturns_to_go=returns_to_go\n",
        "\t\t\t\t\t\t\t\t\t\t\t\t\t)\n",
        "\n",
        "\t\t# only consider non padded elements\n",
        "\t\taction_preds = action_preds.view(-1, act_dim)[traj_mask.view(-1,) > 0]\n",
        "\t\taction_target = action_target.view(-1, act_dim)[traj_mask.view(-1,) > 0]\n",
        "\n",
        "\t\taction_loss = F.mse_loss(action_preds, action_target, reduction='mean')\n",
        "\n",
        "\t\toptimizer.zero_grad()\n",
        "\t\taction_loss.backward()\n",
        "\t\ttorch.nn.utils.clip_grad_norm_(model.parameters(), 0.25)\n",
        "\t\toptimizer.step()\n",
        "\t\tscheduler.step()\n",
        "\n",
        "\t\tlog_action_losses.append(action_loss.detach().cpu().item())\n",
        "\n",
        "\t# evaluate on env\n",
        "\tresults = evaluate_on_env(model, device, context_len, env, rtg_target, rtg_scale,\n",
        "\t                        num_eval_ep, max_eval_ep_len, state_mean, state_std, \n",
        "\t\t\t\t\t\t\t)\n",
        "\teval_avg_reward = results['eval/avg_reward']\n",
        "\teval_avg_ep_len = results['eval/avg_ep_len']\n",
        "\teval_d4rl_score = get_d4rl_normalized_score(results['eval/avg_reward'], env_name) * 100\n",
        "\n",
        "\tmean_action_loss = np.mean(log_action_losses)\n",
        "\ttime_elapsed = str(datetime.now().replace(microsecond=0) - start_time)\n",
        "\n",
        "\ttotal_updates += num_updates_per_iter\n",
        "\n",
        "\tlog_str = (\"=\" * 60 + '\\n' +\n",
        "\t\t\t\"time elapsed: \" + time_elapsed  + '\\n' +\n",
        "\t\t\t\"num of updates: \" + str(total_updates) + '\\n' +\n",
        "\t\t\t\"action loss: \" +  format(mean_action_loss, \".5f\") + '\\n' +\n",
        "\t\t\t\"eval avg reward: \" + format(eval_avg_reward, \".5f\") + '\\n' +\n",
        "\t\t\t\"eval avg ep len: \" + format(eval_avg_ep_len, \".5f\") + '\\n' +\n",
        "\t\t\t\"eval d4rl score: \" + format(eval_d4rl_score, \".5f\")\n",
        "\t\t\t)\n",
        "\n",
        "\tprint(log_str)\n",
        "\n",
        "\tlog_data = [time_elapsed, total_updates, mean_action_loss,\n",
        "\t\t\t\teval_avg_reward, eval_avg_ep_len,\n",
        "\t\t\t\teval_d4rl_score]\n",
        "\n",
        "\tcsv_writer.writerow(log_data)\n",
        "\t\n",
        "\t# save model\n",
        "\tprint(\"max d4rl score: \" + format(max_d4rl_score, \".5f\"))\n",
        "\tif eval_d4rl_score >= max_d4rl_score:\n",
        "\t\tprint(\"saving max d4rl score model at: \" + save_best_model_path)\n",
        "\t\ttorch.save(model.state_dict(), save_best_model_path)\n",
        "\t\tmax_d4rl_score = eval_d4rl_score\n",
        "\n",
        "\tprint(\"saving current model at: \" + save_model_path)\n",
        "\ttorch.save(model.state_dict(), save_model_path)\n",
        "\n",
        "\n",
        "print(\"=\" * 60)\n",
        "print(\"finished training!\")\n",
        "print(\"=\" * 60)\n",
        "end_time = datetime.now().replace(microsecond=0)\n",
        "time_elapsed = str(end_time - start_time)\n",
        "end_time_str = end_time.strftime(\"%y-%m-%d-%H-%M-%S\")\n",
        "print(\"started training at: \" + start_time_str)\n",
        "print(\"finished training at: \" + end_time_str)\n",
        "print(\"total training time: \" + time_elapsed)\n",
        "print(\"max d4rl score: \" + format(max_d4rl_score, \".5f\"))\n",
        "print(\"saved max d4rl score model at: \" + save_best_model_path)\n",
        "print(\"saved last updated model at: \" + save_model_path)\n",
        "print(\"=\" * 60)\n",
        "\n",
        "csv_writer.close()\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "RBLRM5nOVR_8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "prum4oAGlb5P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# test"
      ],
      "metadata": {
        "id": "eosqWqRRJLsZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# set mujoco env path if not already set\n",
        "%env LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.mujoco/mujoco200/bin\n",
        "\n",
        "\n",
        "eval_dataset = \"medium\"\t\t# medium / medium-replay / medium-expert\n",
        "eval_rtg_scale = 1000\t\t# normalize returns to go\n",
        "\n",
        "eval_env_name = \"Walker2d-v3\"\n",
        "eval_rtg_target = 5000\n",
        "eval_env_d4rl_name = f'walker2d-{eval_dataset}-v2'\n",
        "\n",
        "# eval_env_name = \"HalfCheetah-v3\"\n",
        "# eval_rtg_target = 6000\n",
        "# eval_env_d4rl_name = f'halfcheetah-{eval_dataset}-v2'\n",
        "\n",
        "# eval_env_name = \"Hopper-v3\"\n",
        "# eval_rtg_target = 3600\n",
        "# eval_env_d4rl_name = f'hopper-{eval_dataset}-v2'\n",
        "\n",
        "\n",
        "num_test_eval_ep = 10\t\t\t# num of evaluation episodes\n",
        "eval_max_eval_ep_len = 1000\t\t# max len of one episode\n",
        "\n",
        "\n",
        "context_len = 20        # K in decision transformer\n",
        "n_blocks = 3            # num of transformer blocks\n",
        "embed_dim = 128         # embedding (hidden) dim of transformer\n",
        "n_heads = 1             # num of transformer heads\n",
        "dropout_p = 0.1         # dropout probability\n",
        "\n",
        "\n",
        "eval_chk_pt_dir = \"./dt_runs/\"\n",
        "\n",
        "\n",
        "eval_chk_pt_name = \"dt_walker2d-medium-v2_model_22-02-22-09-24-12_best.pt\"\n",
        "eval_chk_pt_list = [eval_chk_pt_name]\n",
        "\n",
        "\n",
        "## manually override check point list\n",
        "## passing a list will evaluate on all checkpoints\n",
        "## and output mean and std score\n",
        "\n",
        "\n",
        "# eval_chk_pt_list = [\n",
        "# \t\"dt_walker2d-medium-v2_model_22-02-20-06-27-12_best.pt\",\n",
        "# \t\"dt_walker2d-medium-v2_model_22-02-20-09-11-30_best.pt\",\n",
        "# \t\"dt_walker2d-medium-v2_model_22-02-22-09-24-12_best.pt\"\n",
        "# ]\n",
        "\n",
        "\n",
        "\n",
        "env_data_stats = get_d4rl_dataset_stats(eval_env_d4rl_name)\n",
        "eval_state_mean = np.array(env_data_stats['state_mean'])\n",
        "eval_state_std = np.array(env_data_stats['state_std'])\n",
        "\n",
        "eval_env = gym.make(eval_env_name)\n",
        "\n",
        "state_dim = eval_env.observation_space.shape[0]\n",
        "act_dim = eval_env.action_space.shape[0]\n",
        "\n",
        "all_scores = []\n",
        "\n",
        "for eval_chk_pt_name in eval_chk_pt_list:\n",
        "\n",
        "\teval_model = DecisionTransformer(\n",
        "\t\t\t\tstate_dim=state_dim,\n",
        "\t\t\t\tact_dim=act_dim,\n",
        "\t\t\t\tn_blocks=n_blocks,\n",
        "\t\t\t\th_dim=embed_dim,\n",
        "\t\t\t\tcontext_len=context_len,\n",
        "\t\t\t\tn_heads=n_heads,\n",
        "\t\t\t\tdrop_p=dropout_p,\n",
        "\t\t\t).to(device)\n",
        "\n",
        "\n",
        "\teval_chk_pt_path = os.path.join(eval_chk_pt_dir, eval_chk_pt_name)\n",
        "\n",
        "\t# load checkpoint\n",
        "\teval_model.load_state_dict(torch.load(eval_chk_pt_path, map_location=device))\n",
        "\n",
        "\tprint(\"model loaded from: \" + eval_chk_pt_path)\n",
        "\n",
        "\t# evaluate on env\n",
        "\tresults = evaluate_on_env(eval_model, device, context_len,\n",
        "\t\t\t\t\t\t\teval_env, eval_rtg_target, eval_rtg_scale,\n",
        "\t\t\t\t\t\t\tnum_test_eval_ep, eval_max_eval_ep_len,\n",
        "\t\t\t\t\t\t\teval_state_mean, eval_state_std)\n",
        "\tprint(results)\n",
        "\n",
        "\tnorm_score = get_d4rl_normalized_score(results['eval/avg_reward'], eval_env_name) * 100\n",
        "\tprint(\"normalized d4rl score: \", norm_score)\n",
        "\n",
        "\tall_scores.append(norm_score)\n",
        "\n",
        "print(\"=\" * 60)\n",
        "all_scores = np.array(all_scores)\n",
        "print(\"evaluated on env: \" + eval_env_name)\n",
        "print(\"total num of checkpoints evaluated: \" + str(len(eval_chk_pt_list)))\n",
        "print(\"d4rl score mean: \" + format(all_scores.mean(), \".5f\"))\n",
        "print(\"d4rl score std: \" + format(all_scores.std(), \".5f\"))\n",
        "print(\"d4rl score var: \" + format(all_scores.var(), \".5f\"))\n",
        "print(\"=\" * 60)\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "Q4-WPlC1VR3q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "SaacQSsJJKVx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## render env\n",
        "\n",
        "\n",
        "\n",
        "*   saves mp4 video of env frames and plays it in notebook\n"
      ],
      "metadata": {
        "id": "wxcJqnb1Him4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "!apt-get install -y xvfb python-opengl ffmpeg > /dev/null 2>&1\n",
        "!pip install -U colabgymrender\n"
      ],
      "metadata": {
        "id": "Hdf_bea2hiRs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "from colabgymrender.recorder import Recorder\n",
        "\n",
        "num_test_eval_ep = 1\n",
        "eval_max_ep_len = 1000\n",
        "\n",
        "\n",
        "directory = \"./render_video\"\n",
        "eval_env = Recorder(eval_env, directory)\n",
        "\n",
        "results = evaluate_on_env(eval_model, device, context_len, \n",
        "                        eval_env, eval_rtg_target, eval_rtg_scale, \n",
        "                        num_test_eval_ep, eval_max_ep_len,\n",
        "\t\t\t\t\t\teval_state_mean, eval_state_std)\n",
        "print(results)\n",
        "\n",
        "norm_score = get_d4rl_normalized_score(results['eval/avg_reward'], eval_env_name) * 100\n",
        "print(\"normalized d4rl score: \", norm_score)\n",
        "\n",
        "eval_env.play()\n",
        "\n",
        "eval_env.close()\n",
        "\n"
      ],
      "metadata": {
        "id": "d-RpNwa0hiPR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "XZNV_H78kRSL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# plots"
      ],
      "metadata": {
        "id": "DjBsdz9mKbZg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import pickle\n",
        "\n",
        "import pandas as pd\n",
        "import glob\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "env_d4rl_name = 'walker2d-medium-v2'\n",
        "\n",
        "log_dir = 'dt_runs/'\n",
        "\n",
        "x_key = \"num_updates\"\n",
        "y_key = \"eval_d4rl_score\"\n",
        "y_smoothing_win = 5\n",
        "plot_avg = False\n",
        "save_fig = False\n",
        "\n",
        "if plot_avg:\n",
        "    save_fig_path = env_d4rl_name + \"_avg.png\"\n",
        "else:\n",
        "    save_fig_path = env_d4rl_name + \".png\"\n",
        "\n",
        "\n",
        "all_files = glob.glob(log_dir + f'/dt_{env_d4rl_name}*.csv')\n",
        "\n",
        "ax = plt.gca()\n",
        "ax.set_title(env_d4rl_name)\n",
        "\n",
        "if plot_avg:\n",
        "    name_list = []\n",
        "    df_list = []\n",
        "    for filename in all_files:\n",
        "        frame = pd.read_csv(filename, index_col=None, header=0)\n",
        "        print(filename, frame.shape)\n",
        "        frame['y_smooth'] = frame[y_key].rolling(window=y_smoothing_win).mean() \n",
        "        df_list.append(frame)\n",
        "    \n",
        "    \n",
        "    df_concat = pd.concat(df_list)\n",
        "    df_concat_groupby = df_concat.groupby(df_concat.index)\n",
        "    data_avg = df_concat_groupby.mean()\n",
        "\n",
        "    data_avg.plot(x=x_key, y='y_smooth', ax=ax)\n",
        "    \n",
        "    ax.set_xlabel(x_key)\n",
        "    ax.set_ylabel(y_key)\n",
        "    ax.legend(['avg of all runs'], loc='lower right')\n",
        "    \n",
        "    if save_fig:\n",
        "        plt.savefig(save_fig_path)\n",
        "        \n",
        "    plt.show()\n",
        "    \n",
        "    \n",
        "else:\n",
        "    name_list = []\n",
        "    for filename in all_files:\n",
        "        frame = pd.read_csv(filename, index_col=None, header=0)\n",
        "        print(filename, frame.shape)\n",
        "        frame['y_smooth'] = frame[y_key].rolling(window=y_smoothing_win).mean()\n",
        "        frame.plot(x=x_key, y='y_smooth', ax=ax)\n",
        "        name_list.append(filename.split('/')[-1])\n",
        "    \n",
        "    ax.set_xlabel(x_key)\n",
        "    ax.set_ylabel(y_key)\n",
        "    ax.legend(name_list, loc='lower right')\n",
        "    \n",
        "    if save_fig:\n",
        "        plt.savefig(save_fig_path)\n",
        "    \n",
        "    plt.show()\n",
        "    \n"
      ],
      "metadata": {
        "id": "2WM69ti2KaRN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "UvUFT19EKaNn"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}