HuggingFaceDecisionTransformer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gJX9tITMv9pT"
      },
      "source": [
        "\n",
        "# Using the HuggingFaceDecision Transformer model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ujAKW4ixZH0X"
      },
      "source": [
        "\n",
        "\n",
        "## Installing dependencies\n",
        "We first need to install the required dependencies in order to run a Mujoco environment inside a Colab notebook\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s7PPg0VZG5cK"
      },
      "outputs": [],
      "source": [
        "!apt-get install -y \\\n",
        "    libgl1-mesa-dev \\\n",
        "    libgl1-mesa-glx \\\n",
        "    libglew-dev \\\n",
        "    libosmesa6-dev \\\n",
        "    software-properties-common \\\n",
        "    patchelf \\\n",
        "    xvfb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZowWDnWoZbtX"
      },
      "source": [
        "## Installing Pip packages\n",
        "We also require the following pip packages:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5w7tslZCHl0T"
      },
      "outputs": [],
      "source": [
        "!pip install gym==0.22.0\n",
        "!pip install free-mujoco-py\n",
        "!pip install git+https://github.com/huggingface/transformers # This change is not part of the pip package, yet.\n",
        "\n",
        "!pip install colabgymrender==1.0.2\n",
        "!pip install xvfbwrapper\n",
        "!pip install imageio==2.4.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "55mjjYn9Z-pu"
      },
      "source": [
        "## Importing the packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TtWrmOCKHD6I"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import mujoco_py\n",
        "import gym\n",
        "import numpy as np\n",
        "\n",
        "from colabgymrender.recorder import Recorder\n",
        "from transformers import DecisionTransformerModel"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TOWpjojcaEBY"
      },
      "source": [
        "## Defining a function that performs masked autoregressive predictive of actions.\n",
        "\n",
        "The model's prediction is conditioned on sequences of states, actions, time-steps and returns. The action for the current time-step is included as zeros and masked in to not skew the model's attention distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WVi8blNvRINK"
      },
      "outputs": [],
      "source": [
        "# Function that gets an action from the model using autoregressive prediction with a window of the previous 20 timesteps.\n",
        "def get_action(model, states, actions, rewards, returns_to_go, timesteps):\n",
        "    # This implementation does not condition on past rewards\n",
        "\n",
        "    states = states.reshape(1, -1, model.config.state_dim)\n",
        "    actions = actions.reshape(1, -1, model.config.act_dim)\n",
        "    returns_to_go = returns_to_go.reshape(1, -1, 1)\n",
        "    timesteps = timesteps.reshape(1, -1)\n",
        "\n",
        "    states = states[:, -model.config.max_length :]\n",
        "    actions = actions[:, -model.config.max_length :]\n",
        "    returns_to_go = returns_to_go[:, -model.config.max_length :]\n",
        "    timesteps = timesteps[:, -model.config.max_length :]\n",
        "    padding = model.config.max_length - states.shape[1]\n",
        "    # pad all tokens to sequence length\n",
        "    attention_mask = torch.cat([torch.zeros(padding), torch.ones(states.shape[1])])\n",
        "    attention_mask = attention_mask.to(dtype=torch.long).reshape(1, -1)\n",
        "    states = torch.cat([torch.zeros((1, padding, model.config.state_dim)), states], dim=1).float()\n",
        "    actions = torch.cat([torch.zeros((1, padding, model.config.act_dim)), actions], dim=1).float()\n",
        "    returns_to_go = torch.cat([torch.zeros((1, padding, 1)), returns_to_go], dim=1).float()\n",
        "    timesteps = torch.cat([torch.zeros((1, padding), dtype=torch.long), timesteps], dim=1)\n",
        "\n",
        "    state_preds, action_preds, return_preds = model(\n",
        "        states=states,\n",
        "        actions=actions,\n",
        "        rewards=rewards,\n",
        "        returns_to_go=returns_to_go,\n",
        "        timesteps=timesteps,\n",
        "        attention_mask=attention_mask,\n",
        "        return_dict=False,\n",
        "    )\n",
        "\n",
        "    return action_preds[0, -1]"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Hopper"
      ],
      "metadata": {
        "id": "s8Y9FFxPW8Gf"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2PW-G_nvUIOd"
      },
      "outputs": [],
      "source": [
        "# build the environment\n",
        "directory = './video'\n",
        "\n",
        "env = gym.make(\"Hopper-v3\")\n",
        "env = Recorder(env, directory, fps=30)\n",
        "max_ep_len = 1000\n",
        "device = \"cpu\"\n",
        "scale = 1000.0  # normalization for rewards/returns\n",
        "TARGET_RETURN = 3600 / scale  # evaluation is conditioned on a return of 3600, scaled accordingly\n",
        "# mean and std computed from training dataset these are available in the model card for each model.\n",
        "# Expert\n",
        "state_mean = np.array(\n",
        "    [1.3490015,  -0.11208222, -0.5506444,  -0.13188992, -0.00378754,  2.6071432,\n",
        "     0.02322114, -0.01626922, -0.06840388, -0.05183131,  0.04272673,]\n",
        ")\n",
        "state_std = np.array(\n",
        "    [0.15980862, 0.0446214,  0.14307782, 0.17629202, 0.5912333,  0.5899924,\n",
        " 1.5405099,  0.8152689,  2.0173461,  2.4107876,  5.8440027,]\n",
        ")\n",
        "\n",
        "# # Medium\n",
        "# state_mean = np.array(\n",
        "#     [ 1.311279, -0.08469521, -0.5382719, -0.07201576, 0.04932366, 2.1066856,\n",
        "#      -0.15017354, 0.00878345, -0.2848186, -0.18540096, -0.28461286]\n",
        "# )\n",
        "# state_std = np.array(\n",
        "#     [0.17790751, 0.05444621, 0.21297139, 0.14530419, 0.6124444, 0.85174465,\n",
        "#      1.4515252, 0.6751696, 1.536239, 1.6160746, 5.6072536 ]\n",
        "# )\n",
        "\n",
        "state_dim = env.observation_space.shape[0]\n",
        "act_dim = env.action_space.shape[0]\n",
        "# Create the decision transformer model\n",
        "model = DecisionTransformerModel.from_pretrained(\"edbeeching/decision-transformer-gym-hopper-expert\")\n",
        "model = model.to(device)\n",
        "print(list(model.encoder.wpe.parameters()))\n",
        "\n",
        "state_mean = torch.from_numpy(state_mean).to(device=device)\n",
        "state_std = torch.from_numpy(state_std).to(device=device)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Interact with the environment and create a video\n",
        "r = []\n",
        "l = []\n",
        "for i in range(5):\n",
        "    episode_return, episode_length = 0, 0\n",
        "    state = env.reset()\n",
        "    target_return = torch.tensor(TARGET_RETURN, device=device, dtype=torch.float32).reshape(1, 1)\n",
        "    states = torch.from_numpy(state).reshape(1, state_dim).to(device=device, dtype=torch.float32)\n",
        "    actions = torch.zeros((0, act_dim), device=device, dtype=torch.float32)\n",
        "    rewards = torch.zeros(0, device=device, dtype=torch.float32)\n",
        "\n",
        "    timesteps = torch.tensor(0, device=device, dtype=torch.long).reshape(1, 1)\n",
        "    for t in range(max_ep_len):\n",
        "        actions = torch.cat([actions, torch.zeros((1, act_dim), device=device)], dim=0)\n",
        "        rewards = torch.cat([rewards, torch.zeros(1, device=device)])\n",
        "\n",
        "        action = get_action(\n",
        "            model,\n",
        "            (states - state_mean) / state_std,\n",
        "            actions,\n",
        "            rewards,\n",
        "            target_return,\n",
        "            timesteps,\n",
        "        )\n",
        "        actions[-1] = action\n",
        "        action = action.detach().cpu().numpy()\n",
        "\n",
        "        state, reward, done, _ = env.step(action)\n",
        "\n",
        "        cur_state = torch.from_numpy(state).to(device=device).reshape(1, state_dim)\n",
        "        states = torch.cat([states, cur_state], dim=0)\n",
        "        rewards[-1] = reward\n",
        "\n",
        "        pred_return = target_return[0, -1] - (reward / scale)\n",
        "        target_return = torch.cat([target_return, pred_return.reshape(1, 1)], dim=1)\n",
        "        timesteps = torch.cat([timesteps, torch.ones((1, 1), device=device, dtype=torch.long) * (t + 1)], dim=1)\n",
        "\n",
        "        episode_return += reward\n",
        "        episode_length += 1\n",
        "\n",
        "        if done:\n",
        "            break\n",
        "    r.append(episode_return)\n",
        "    l.append(episode_length)\n",
        "    print(episode_return, episode_length)"
      ],
      "metadata": {
        "id": "6W9OsPS1T5pQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Mean and Std. of Return and Episode Length\n",
        "m_r = np.mean(r)\n",
        "s_r = np.std(r)\n",
        "m_l = np.mean(l)\n",
        "s_l = np.std(l)\n",
        "print(m_r, s_r)\n",
        "print(m_l, s_l)"
      ],
      "metadata": {
        "id": "6RKOAdS2Tsvn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Play the video\n",
        "env.play()\n",
        "# If you want to convert the video:\n",
        "# !ffmpeg -i {your_video} -vcodec h264 replay.mp4"
      ],
      "metadata": {
        "id": "FC5AhsK6UDff"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Walker 2D"
      ],
      "metadata": {
        "id": "Hw4bxU6FT6k3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# build the environment\n",
        "directory = './video'\n",
        "\n",
        "env = gym.make(\"Walker2d-v3\")\n",
        "env = Recorder(env, directory, fps=30)\n",
        "max_ep_len = 1000\n",
        "device = \"cpu\"\n",
        "scale = 1000.0  # normalization for rewards/returns\n",
        "TARGET_RETURN = 5000 / scale  # evaluation is conditioned on a return of 3600, scaled accordingly\n",
        "# mean and std computed from training dataset these are available in the model card for each model.\n",
        "# Expert\n",
        "state_mean = np.array(\n",
        "     [ 1.218966, 0.14163373, -0.03704914, -0.1381431, 0.51382244, -0.0471911,\n",
        "      -0.47288352, 0.04225416, 2.3948874, -0.03143199, 0.04466356, -0.02390724,\n",
        "       -0.10134014, 0.09090938, -0.00419264, -0.12120572, -0.5497064]\n",
        ")\n",
        "state_std = np.array(\n",
        "    [0.12311358, 0.324188, 0.11456084, 0.26230657, 0.5640279, 0.22718786,\n",
        "     0.38373196, 0.7373677, 1.2387927, 0.7980206, 1.5664079, 1.8092705,\n",
        "     3.0256042, 4.062486, 1.4586568, 3.744569, 5.585129 ]\n",
        ")\n",
        "\n",
        "# # Medium\n",
        "# state_mean = np.array(\n",
        "#       [ 1.218966, 0.14163373, -0.03704914, -0.1381431, 0.51382244, -0.0471911,\n",
        "#        -0.47288352, 0.04225416, 2.3948874, -0.03143199, 0.04466356, -0.02390724,\n",
        "#         -0.10134014, 0.09090938, -0.00419264, -0.12120572, -0.5497064]\n",
        "# )\n",
        "# state_std = np.array(\n",
        "#     [0.12311358, 0.324188, 0.11456084, 0.26230657, 0.5640279, 0.22718786,\n",
        "#      0.38373196, 0.7373677, 1.2387927, 0.7980206, 1.5664079, 1.8092705,\n",
        "#      3.0256042, 4.062486, 1.4586568, 3.744569, 5.585129 ]\n",
        "# )\n",
        "state_dim = env.observation_space.shape[0]\n",
        "act_dim = env.action_space.shape[0]\n",
        "# Create the decision transformer model\n",
        "model = DecisionTransformerModel.from_pretrained(\"edbeeching/decision-transformer-gym-walker2d-expert\")\n",
        "model = model.to(device)\n",
        "print(list(model.encoder.wpe.parameters()))\n",
        "\n",
        "state_mean = torch.from_numpy(state_mean).to(device=device)\n",
        "state_std = torch.from_numpy(state_std).to(device=device)"
      ],
      "metadata": {
        "id": "NySBr6n2qeBZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FuM7LeZ4Veh0"
      },
      "outputs": [],
      "source": [
        "# Interact with the environment and create a video\n",
        "r = []\n",
        "l = []\n",
        "for i in range(5):\n",
        "    episode_return, episode_length = 0, 0\n",
        "    state = env.reset()\n",
        "    target_return = torch.tensor(TARGET_RETURN, device=device, dtype=torch.float32).reshape(1, 1)\n",
        "    states = torch.from_numpy(state).reshape(1, state_dim).to(device=device, dtype=torch.float32)\n",
        "    actions = torch.zeros((0, act_dim), device=device, dtype=torch.float32)\n",
        "    rewards = torch.zeros(0, device=device, dtype=torch.float32)\n",
        "\n",
        "    timesteps = torch.tensor(0, device=device, dtype=torch.long).reshape(1, 1)\n",
        "    for t in range(max_ep_len):\n",
        "        actions = torch.cat([actions, torch.zeros((1, act_dim), device=device)], dim=0)\n",
        "        rewards = torch.cat([rewards, torch.zeros(1, device=device)])\n",
        "\n",
        "        action = get_action(\n",
        "            model,\n",
        "            (states - state_mean) / state_std,\n",
        "            actions,\n",
        "            rewards,\n",
        "            target_return,\n",
        "            timesteps,\n",
        "        )\n",
        "        actions[-1] = action\n",
        "        action = action.detach().cpu().numpy()\n",
        "\n",
        "        state, reward, done, _ = env.step(action)\n",
        "\n",
        "        cur_state = torch.from_numpy(state).to(device=device).reshape(1, state_dim)\n",
        "        states = torch.cat([states, cur_state], dim=0)\n",
        "        rewards[-1] = reward\n",
        "\n",
        "        pred_return = target_return[0, -1] - (reward / scale)\n",
        "        target_return = torch.cat([target_return, pred_return.reshape(1, 1)], dim=1)\n",
        "        timesteps = torch.cat([timesteps, torch.ones((1, 1), device=device, dtype=torch.long) * (t + 1)], dim=1)\n",
        "\n",
        "        episode_return += reward\n",
        "        episode_length += 1\n",
        "\n",
        "        if done:\n",
        "            break\n",
        "    r.append(episode_return)\n",
        "    l.append(episode_length)\n",
        "    print(episode_return, episode_length)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Mean and Std. of Return and Episode Length\n",
        "m_r = np.mean(r)\n",
        "s_r = np.std(r)\n",
        "m_l = np.mean(l)\n",
        "s_l = np.std(l)\n",
        "print(m_r, s_r)\n",
        "print(m_l, s_l)"
      ],
      "metadata": {
        "id": "ZuQdzefJWq24"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2l1bAG8n1Z64"
      },
      "outputs": [],
      "source": [
        "# Play the video\n",
        "env.play()\n",
        "# If you want to convert the video:\n",
        "# !ffmpeg -i {your_video} -vcodec h264 replay.mp4"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Half Cheetah"
      ],
      "metadata": {
        "id": "LCLQ1ipMUH_C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# build the environment\n",
        "directory = './video'\n",
        "\n",
        "env = gym.make(\"HalfCheetah-v3\")\n",
        "env = Recorder(env, directory, fps=30)\n",
        "max_ep_len = 1000\n",
        "device = \"cpu\"\n",
        "scale = 1000.0  # normalization for rewards/returns\n",
        "TARGET_RETURN = 6000 / scale  # evaluation is conditioned on a return of 3600, scaled accordingly\n",
        "# mean and std computed from training dataset these are available in the model card for each model.\n",
        "# Expert\n",
        "state_mean = np.array(\n",
        "     [ -0.04489148, 0.03232588, 0.06034835, -0.17081226, -0.19480659, -0.05751596,\n",
        "      0.09701628, 0.03239211, 11.047426, -0.07997331, -0.32363534, 0.36297753,\n",
        "       0.42322603, 0.40836546, 1.1085187, -0.4874403, -0.0737481 ]\n",
        ")\n",
        "state_std = np.array(\n",
        "    [0.04002118, 0.4107858, 0.54217845, 0.41522816, 0.23796624, 0.62036866,\n",
        "     0.30100912, 0.21737163, 2.2105937, 0.572586, 1.7255033, 11.844218,\n",
        "     12.06324, 7.0495934, 13.499867, 7.195647, 5.0264325]\n",
        ")\n",
        "\n",
        "# # Medium\n",
        "# state_mean = np.array(\n",
        "#      [-0.06845774, 0.01641455, -0.18354906, -0.27624607, -0.34061527, -0.09339716,\n",
        "#       -0.21321271, -0.08774239, 5.1730075, -0.04275195, -0.03610836, 0.14053793,\n",
        "#       0.06049833, 0.09550975, 0.067391, 0.00562739, 0.01338279]\n",
        "# )\n",
        "# state_std = np.array(\n",
        "#     [0.07472999, 0.30234998, 0.3020731, 0.34417078, 0.17619242, 0.5072056,\n",
        "#      0.25670078, 0.32948127, 1.2574149, 0.7600542, 1.9800916, 6.5653625,\n",
        "#      7.4663677, 4.472223, 10.566964, 5.6719327, 7.498259]\n",
        "# )\n",
        "state_dim = env.observation_space.shape[0]\n",
        "act_dim = env.action_space.shape[0]\n",
        "# Create the decision transformer model\n",
        "model = DecisionTransformerModel.from_pretrained(\"edbeeching/decision-transformer-gym-halfcheetah-expert\")\n",
        "model = model.to(device)\n",
        "print(list(model.encoder.wpe.parameters()))\n",
        "\n",
        "state_mean = torch.from_numpy(state_mean).to(device=device)\n",
        "state_std = torch.from_numpy(state_std).to(device=device)"
      ],
      "metadata": {
        "id": "MrS4rr1EUHRS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Interact with the environment and create a video\n",
        "r = []\n",
        "l = []\n",
        "for i in range(5):\n",
        "    episode_return, episode_length = 0, 0\n",
        "    state = env.reset()\n",
        "    target_return = torch.tensor(TARGET_RETURN, device=device, dtype=torch.float32).reshape(1, 1)\n",
        "    states = torch.from_numpy(state).reshape(1, state_dim).to(device=device, dtype=torch.float32)\n",
        "    actions = torch.zeros((0, act_dim), device=device, dtype=torch.float32)\n",
        "    rewards = torch.zeros(0, device=device, dtype=torch.float32)\n",
        "\n",
        "    timesteps = torch.tensor(0, device=device, dtype=torch.long).reshape(1, 1)\n",
        "    for t in range(max_ep_len):\n",
        "        actions = torch.cat([actions, torch.zeros((1, act_dim), device=device)], dim=0)\n",
        "        rewards = torch.cat([rewards, torch.zeros(1, device=device)])\n",
        "\n",
        "        action = get_action(\n",
        "            model,\n",
        "            (states - state_mean) / state_std,\n",
        "            actions,\n",
        "            rewards,\n",
        "            target_return,\n",
        "            timesteps,\n",
        "        )\n",
        "        actions[-1] = action\n",
        "        action = action.detach().cpu().numpy()\n",
        "\n",
        "        state, reward, done, _ = env.step(action)\n",
        "\n",
        "        cur_state = torch.from_numpy(state).to(device=device).reshape(1, state_dim)\n",
        "        states = torch.cat([states, cur_state], dim=0)\n",
        "        rewards[-1] = reward\n",
        "\n",
        "        pred_return = target_return[0, -1] - (reward / scale)\n",
        "        target_return = torch.cat([target_return, pred_return.reshape(1, 1)], dim=1)\n",
        "        timesteps = torch.cat([timesteps, torch.ones((1, 1), device=device, dtype=torch.long) * (t + 1)], dim=1)\n",
        "\n",
        "        episode_return += reward\n",
        "        episode_length += 1\n",
        "\n",
        "        if done:\n",
        "            break\n",
        "    r.append(episode_return)\n",
        "    l.append(episode_length)\n",
        "    print(episode_return, episode_length)"
      ],
      "metadata": {
        "id": "sqDLzUVuo9Sq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Mean and Std. of Return and Episode Length\n",
        "m_r = np.mean(r)\n",
        "s_r = np.std(r)\n",
        "m_l = np.mean(l)\n",
        "s_l = np.std(l)\n",
        "print(m_r, s_r)\n",
        "print(m_l, s_l)"
      ],
      "metadata": {
        "id": "jEBx2vEpWvgN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Play the video\n",
        "env.play()\n",
        "# If you want to convert the video:\n",
        "# !ffmpeg -i {your_video} -vcodec h264 replay.mp4"
      ],
      "metadata": {
        "id": "A9s4yIunUQnK"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}