TD3_Half_Humanoid.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TD3_Half_Humanoid.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WXu1r8qvSzWf",
        "colab_type": "text"
      },
      "source": [
        "# Twin-Delayed DDPG"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YRzQUhuUTc0J",
        "colab_type": "text"
      },
      "source": [
        "## Installing the packages"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HAHMB0Ze8fU0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install pybullet"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xjm2onHdT-Av",
        "colab_type": "text"
      },
      "source": [
        "## Importing the libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ikr2p0Js8iB4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import time\n",
        "import random\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import pybullet_envs\n",
        "import gym\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from gym import wrappers\n",
        "from torch.autograd import Variable\n",
        "from collections import deque"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y2nGdtlKVydr",
        "colab_type": "text"
      },
      "source": [
        "## Step 1: We initialize the Experience Replay memory"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u5rW0IDB8nTO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ReplayBuffer(object):\n",
        "\n",
        "  def __init__(self, max_size=1e6):\n",
        "    self.storage = []\n",
        "    self.max_size = max_size\n",
        "    self.ptr = 0\n",
        "\n",
        "  def add(self, transition):\n",
        "    if len(self.storage) == self.max_size:\n",
        "      self.storage[int(self.ptr)] = transition\n",
        "      self.ptr = (self.ptr + 1) % self.max_size\n",
        "    else:\n",
        "      self.storage.append(transition)\n",
        "\n",
        "  def sample(self, batch_size):\n",
        "    ind = np.random.randint(0, len(self.storage), size=batch_size)\n",
        "    batch_states, batch_next_states, batch_actions, batch_rewards, batch_dones = [], [], [], [], []\n",
        "    for i in ind: \n",
        "      state, next_state, action, reward, done = self.storage[i]\n",
        "      batch_states.append(np.array(state, copy=False))\n",
        "      batch_next_states.append(np.array(next_state, copy=False))\n",
        "      batch_actions.append(np.array(action, copy=False))\n",
        "      batch_rewards.append(np.array(reward, copy=False))\n",
        "      batch_dones.append(np.array(done, copy=False))\n",
        "    return np.array(batch_states), np.array(batch_next_states), np.array(batch_actions), np.array(batch_rewards).reshape(-1, 1), np.array(batch_dones).reshape(-1, 1)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jb7TTaHxWbQD",
        "colab_type": "text"
      },
      "source": [
        "## Step 2: We build one neural network for the Actor model and one neural network for the Actor target"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4CeRW4D79HL0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Actor(nn.Module):\n",
        "  \n",
        "  def __init__(self, state_dim, action_dim, max_action):\n",
        "    super(Actor, self).__init__()\n",
        "    self.layer_1 = nn.Linear(state_dim, 400)\n",
        "    self.layer_2 = nn.Linear(400, 300)\n",
        "    self.layer_3 = nn.Linear(300, action_dim)\n",
        "    self.max_action = max_action\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = F.relu(self.layer_1(x))\n",
        "    x = F.relu(self.layer_2(x))\n",
        "    x = self.max_action * torch.tanh(self.layer_3(x))\n",
        "    return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HRDDce8FXef7",
        "colab_type": "text"
      },
      "source": [
        "## Step 3: We build two neural networks for the two Critic models and two neural networks for the two Critic targets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OCee7gwR9Jrs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Critic(nn.Module):\n",
        "  \n",
        "  def __init__(self, state_dim, action_dim):\n",
        "    super(Critic, self).__init__()\n",
        "    # Defining the first Critic neural network\n",
        "    self.layer_1 = nn.Linear(state_dim + action_dim, 400)\n",
        "    self.layer_2 = nn.Linear(400, 300)\n",
        "    self.layer_3 = nn.Linear(300, 1)\n",
        "    # Defining the second Critic neural network\n",
        "    self.layer_4 = nn.Linear(state_dim + action_dim, 400)\n",
        "    self.layer_5 = nn.Linear(400, 300)\n",
        "    self.layer_6 = nn.Linear(300, 1)\n",
        "\n",
        "  def forward(self, x, u):\n",
        "    xu = torch.cat([x, u], 1)\n",
        "    # Forward-Propagation on the first Critic Neural Network\n",
        "    x1 = F.relu(self.layer_1(xu))\n",
        "    x1 = F.relu(self.layer_2(x1))\n",
        "    x1 = self.layer_3(x1)\n",
        "    # Forward-Propagation on the second Critic Neural Network\n",
        "    x2 = F.relu(self.layer_4(xu))\n",
        "    x2 = F.relu(self.layer_5(x2))\n",
        "    x2 = self.layer_6(x2)\n",
        "    return x1, x2\n",
        "\n",
        "  def Q1(self, x, u):\n",
        "    xu = torch.cat([x, u], 1)\n",
        "    x1 = F.relu(self.layer_1(xu))\n",
        "    x1 = F.relu(self.layer_2(x1))\n",
        "    x1 = self.layer_3(x1)\n",
        "    return x1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NzIDuONodenW",
        "colab_type": "text"
      },
      "source": [
        "## Steps 4 to 15: Training Process"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "zzd0H1xukdKe",
        "colab": {}
      },
      "source": [
        "# Selecting the device (CPU or GPU)\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Building the whole Training Process into a class\n",
        "\n",
        "class TD3(object):\n",
        "  \n",
        "  def __init__(self, state_dim, action_dim, max_action):\n",
        "    self.actor = Actor(state_dim, action_dim, max_action).to(device)\n",
        "    self.actor_target = Actor(state_dim, action_dim, max_action).to(device)\n",
        "    self.actor_target.load_state_dict(self.actor.state_dict())\n",
        "    self.actor_optimizer = torch.optim.Adam(self.actor.parameters())\n",
        "    self.critic = Critic(state_dim, action_dim).to(device)\n",
        "    self.critic_target = Critic(state_dim, action_dim).to(device)\n",
        "    self.critic_target.load_state_dict(self.critic.state_dict())\n",
        "    self.critic_optimizer = torch.optim.Adam(self.critic.parameters())\n",
        "    self.max_action = max_action\n",
        "\n",
        "  def select_action(self, state):\n",
        "    state = torch.Tensor(state.reshape(1, -1)).to(device)\n",
        "    return self.actor(state).cpu().data.numpy().flatten()\n",
        "\n",
        "  def train(self, replay_buffer, iterations, batch_size=100, discount=0.99, tau=0.005, policy_noise=0.2, noise_clip=0.5, policy_freq=2):\n",
        "    \n",
        "    for it in range(iterations):\n",
        "      \n",
        "      # Step 4: We sample a batch of transitions (s, s’, a, r) from the memory\n",
        "      batch_states, batch_next_states, batch_actions, batch_rewards, batch_dones = replay_buffer.sample(batch_size)\n",
        "      state = torch.Tensor(batch_states).to(device)\n",
        "      next_state = torch.Tensor(batch_next_states).to(device)\n",
        "      action = torch.Tensor(batch_actions).to(device)\n",
        "      reward = torch.Tensor(batch_rewards).to(device)\n",
        "      done = torch.Tensor(batch_dones).to(device)\n",
        "      \n",
        "      # Step 5: From the next state s’, the Actor target plays the next action a’\n",
        "      next_action = self.actor_target(next_state)\n",
        "      \n",
        "      # Step 6: We add Gaussian noise to this next action a’ and we clamp it in a range of values supported by the environment\n",
        "      noise = torch.Tensor(batch_actions).data.normal_(0, policy_noise).to(device)\n",
        "      noise = noise.clamp(-noise_clip, noise_clip)\n",
        "      next_action = (next_action + noise).clamp(-self.max_action, self.max_action)\n",
        "      \n",
        "      # Step 7: The two Critic targets take each the couple (s’, a’) as input and return two Q-values Qt1(s’,a’) and Qt2(s’,a’) as outputs\n",
        "      target_Q1, target_Q2 = self.critic_target(next_state, next_action)\n",
        "      \n",
        "      # Step 8: We keep the minimum of these two Q-values: min(Qt1, Qt2)\n",
        "      target_Q = torch.min(target_Q1, target_Q2)\n",
        "      \n",
        "      # Step 9: We get the final target of the two Critic models, which is: Qt = r + γ * min(Qt1, Qt2), where γ is the discount factor\n",
        "      target_Q = reward + ((1 - done) * discount * target_Q).detach()\n",
        "      \n",
        "      # Step 10: The two Critic models take each the couple (s, a) as input and return two Q-values Q1(s,a) and Q2(s,a) as outputs\n",
        "      current_Q1, current_Q2 = self.critic(state, action)\n",
        "      \n",
        "      # Step 11: We compute the loss coming from the two Critic models: Critic Loss = MSE_Loss(Q1(s,a), Qt) + MSE_Loss(Q2(s,a), Qt)\n",
        "      critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)\n",
        "      \n",
        "      # Step 12: We backpropagate this Critic loss and update the parameters of the two Critic models with a SGD optimizer\n",
        "      self.critic_optimizer.zero_grad()\n",
        "      critic_loss.backward()\n",
        "      self.critic_optimizer.step()\n",
        "      \n",
        "      # Step 13: Once every two iterations, we update our Actor model by performing gradient ascent on the output of the first Critic model\n",
        "      if it % policy_freq == 0:\n",
        "        actor_loss = -self.critic.Q1(state, self.actor(state)).mean()\n",
        "        self.actor_optimizer.zero_grad()\n",
        "        actor_loss.backward()\n",
        "        self.actor_optimizer.step()\n",
        "        \n",
        "        # Step 14: Still once every two iterations, we update the weights of the Actor target by polyak averaging\n",
        "        for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):\n",
        "          target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)\n",
        "        \n",
        "        # Step 15: Still once every two iterations, we update the weights of the Critic target by polyak averaging\n",
        "        for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):\n",
        "          target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)\n",
        "  \n",
        "  # Making a save method to save a trained model\n",
        "  def save(self, filename, directory):\n",
        "    torch.save(self.actor.state_dict(), '%s/%s_actor.pth' % (directory, filename))\n",
        "    torch.save(self.critic.state_dict(), '%s/%s_critic.pth' % (directory, filename))\n",
        "  \n",
        "  # Making a load method to load a pre-trained model\n",
        "  def load(self, filename, directory):\n",
        "    self.actor.load_state_dict(torch.load('%s/%s_actor.pth' % (directory, filename)))\n",
        "    self.critic.load_state_dict(torch.load('%s/%s_critic.pth' % (directory, filename)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ka-ZRtQvjBex",
        "colab_type": "text"
      },
      "source": [
        "## We make a function that evaluates the policy by calculating its average reward over 10 episodes"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qabqiYdp9wDM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def evaluate_policy(policy, eval_episodes=10):\n",
        "  avg_reward = 0.\n",
        "  for _ in range(eval_episodes):\n",
        "    obs = env.reset()\n",
        "    done = False\n",
        "    while not done:\n",
        "      action = policy.select_action(np.array(obs))\n",
        "      obs, reward, done, _ = env.step(action)\n",
        "      avg_reward += reward\n",
        "  avg_reward /= eval_episodes\n",
        "  print (\"---------------------------------------\")\n",
        "  print (\"Average Reward over the Evaluation Step: %f\" % (avg_reward))\n",
        "  print (\"---------------------------------------\")\n",
        "  return avg_reward"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gGuKmH_ijf7U",
        "colab_type": "text"
      },
      "source": [
        "## We set the parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HFj6wbAo97lk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env_name = \"Walker2DBulletEnv-v0\" # Name of a environment (set it to any Continous environment you want)\n",
        "seed = 0 # Random seed number\n",
        "start_timesteps = 1e4 # Number of iterations/timesteps before which the model randomly chooses an action, and after which it starts to use the policy network\n",
        "eval_freq = 5e3 # How often the evaluation step is performed (after how many timesteps)\n",
        "max_timesteps = 5e5 # Total number of iterations/timesteps\n",
        "save_models = True # Boolean checker whether or not to save the pre-trained model\n",
        "expl_noise = 0.1 # Exploration noise - STD value of exploration Gaussian noise\n",
        "batch_size = 100 # Size of the batch\n",
        "discount = 0.99 # Discount factor gamma, used in the calculation of the total discounted reward\n",
        "tau = 0.005 # Target network update rate\n",
        "policy_noise = 0.2 # STD of Gaussian noise added to the actions for the exploration purposes\n",
        "noise_clip = 0.5 # Maximum value of the Gaussian noise added to the actions (policy)\n",
        "policy_freq = 2 # Number of iterations to wait before the policy network (Actor model) is updated"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hjwf2HCol3XP",
        "colab_type": "text"
      },
      "source": [
        "## We create a file name for the two saved models: the Actor and Critic models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1fyH8N5z-o3o",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "file_name = \"%s_%s_%s\" % (\"TD3\", env_name, str(seed))\n",
        "print (\"---------------------------------------\")\n",
        "print (\"Settings: %s\" % (file_name))\n",
        "print (\"---------------------------------------\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kop-C96Aml8O",
        "colab_type": "text"
      },
      "source": [
        "## We create a folder inside which will be saved the trained models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Src07lvY-zXb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "if not os.path.exists(\"./results\"):\n",
        "  os.makedirs(\"./results\")\n",
        "if save_models and not os.path.exists(\"./pytorch_models\"):\n",
        "  os.makedirs(\"./pytorch_models\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qEAzOd47mv1Z",
        "colab_type": "text"
      },
      "source": [
        "## We create the PyBullet environment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CyQXJUIs-6BV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env = gym.make(env_name)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5YdPG4HXnNsh",
        "colab_type": "text"
      },
      "source": [
        "## We set seeds and we get the necessary information on the states and actions in the chosen environment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z3RufYec_ADj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "np.random.seed(seed)\n",
        "state_dim = env.observation_space.shape[0]\n",
        "action_dim = env.action_space.shape[0]\n",
        "max_action = float(env.action_space.high[0])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HWEgDAQxnbem",
        "colab_type": "text"
      },
      "source": [
        "## We create the policy network (the Actor model)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wTVvG7F8_EWg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "policy = TD3(state_dim, action_dim, max_action)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZI60VN2Unklh",
        "colab_type": "text"
      },
      "source": [
        "## We create the Experience Replay memory"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sd-ZsdXR_LgV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "replay_buffer = ReplayBuffer()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QYOpCyiDnw7s",
        "colab_type": "text"
      },
      "source": [
        "## We define a list where all the evaluation results over 10 episodes are stored"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dhC_5XJ__Orp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "evaluations = [evaluate_policy(policy)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xm-4b3p6rglE",
        "colab_type": "text"
      },
      "source": [
        "## We create a new folder directory in which the final results (videos of the agent) will be populated"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MTL9uMd0ru03",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def mkdir(base, name):\n",
        "    path = os.path.join(base, name)\n",
        "    if not os.path.exists(path):\n",
        "        os.makedirs(path)\n",
        "    return path\n",
        "work_dir = mkdir('exp', 'brs')\n",
        "monitor_dir = mkdir(work_dir, 'monitor')\n",
        "max_episode_steps = env._max_episode_steps\n",
        "save_env_vid = False\n",
        "if save_env_vid:\n",
        "  env = wrappers.Monitor(env, monitor_dir, force = True)\n",
        "  env.reset()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "31n5eb03p-Fm",
        "colab_type": "text"
      },
      "source": [
        "## We initialize the variables"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1vN5EvxK_QhT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "total_timesteps = 0\n",
        "timesteps_since_eval = 0\n",
        "episode_num = 0\n",
        "done = True\n",
        "t0 = time.time()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q9gsjvtPqLgT",
        "colab_type": "text"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y_ouY4NH_Y0I",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# We start the main loop over 500,000 timesteps\n",
        "while total_timesteps < max_timesteps:\n",
        "  \n",
        "  # If the episode is done\n",
        "  if done:\n",
        "\n",
        "    # If we are not at the very beginning, we start the training process of the model\n",
        "    if total_timesteps != 0:\n",
        "      print(\"Total Timesteps: {} Episode Num: {} Reward: {}\".format(total_timesteps, episode_num, episode_reward))\n",
        "      policy.train(replay_buffer, episode_timesteps, batch_size, discount, tau, policy_noise, noise_clip, policy_freq)\n",
        "\n",
        "    # We evaluate the episode and we save the policy\n",
        "    if timesteps_since_eval >= eval_freq:\n",
        "      timesteps_since_eval %= eval_freq\n",
        "      evaluations.append(evaluate_policy(policy))\n",
        "      policy.save(file_name, directory=\"./pytorch_models\")\n",
        "      np.save(\"./results/%s\" % (file_name), evaluations)\n",
        "    \n",
        "    # When the training step is done, we reset the state of the environment\n",
        "    obs = env.reset()\n",
        "    \n",
        "    # Set the Done to False\n",
        "    done = False\n",
        "    \n",
        "    # Set rewards and episode timesteps to zero\n",
        "    episode_reward = 0\n",
        "    episode_timesteps = 0\n",
        "    episode_num += 1\n",
        "  \n",
        "  # Before 10000 timesteps, we play random actions\n",
        "  if total_timesteps < start_timesteps:\n",
        "    action = env.action_space.sample()\n",
        "  else: # After 10000 timesteps, we switch to the model\n",
        "    action = policy.select_action(np.array(obs))\n",
        "    # If the explore_noise parameter is not 0, we add noise to the action and we clip it\n",
        "    if expl_noise != 0:\n",
        "      action = (action + np.random.normal(0, expl_noise, size=env.action_space.shape[0])).clip(env.action_space.low, env.action_space.high)\n",
        "  \n",
        "  # The agent performs the action in the environment, then reaches the next state and receives the reward\n",
        "  new_obs, reward, done, _ = env.step(action)\n",
        "  \n",
        "  # We check if the episode is done\n",
        "  done_bool = 0 if episode_timesteps + 1 == env._max_episode_steps else float(done)\n",
        "  \n",
        "  # We increase the total reward\n",
        "  episode_reward += reward\n",
        "  \n",
        "  # We store the new transition into the Experience Replay memory (ReplayBuffer)\n",
        "  replay_buffer.add((obs, new_obs, action, reward, done_bool))\n",
        "\n",
        "  # We update the state, the episode timestep, the total timesteps, and the timesteps since the evaluation of the policy\n",
        "  obs = new_obs\n",
        "  episode_timesteps += 1\n",
        "  total_timesteps += 1\n",
        "  timesteps_since_eval += 1\n",
        "\n",
        "# We add the last policy evaluation to our list of evaluations and we save our model\n",
        "evaluations.append(evaluate_policy(policy))\n",
        "if save_models: policy.save(\"%s\" % (file_name), directory=\"./pytorch_models\")\n",
        "np.save(\"./results/%s\" % (file_name), evaluations)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wi6e2-_pu05e",
        "colab_type": "text"
      },
      "source": [
        "## Inference"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oW4d1YAMqif1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Actor(nn.Module):\n",
        "  \n",
        "  def __init__(self, state_dim, action_dim, max_action):\n",
        "    super(Actor, self).__init__()\n",
        "    self.layer_1 = nn.Linear(state_dim, 400)\n",
        "    self.layer_2 = nn.Linear(400, 300)\n",
        "    self.layer_3 = nn.Linear(300, action_dim)\n",
        "    self.max_action = max_action\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = F.relu(self.layer_1(x))\n",
        "    x = F.relu(self.layer_2(x))\n",
        "    x = self.max_action * torch.tanh(self.layer_3(x)) \n",
        "    return x\n",
        "\n",
        "class Critic(nn.Module):\n",
        "  \n",
        "  def __init__(self, state_dim, action_dim):\n",
        "    super(Critic, self).__init__()\n",
        "    # Defining the first Critic neural network\n",
        "    self.layer_1 = nn.Linear(state_dim + action_dim, 400)\n",
        "    self.layer_2 = nn.Linear(400, 300)\n",
        "    self.layer_3 = nn.Linear(300, 1)\n",
        "    # Defining the second Critic neural network\n",
        "    self.layer_4 = nn.Linear(state_dim + action_dim, 400)\n",
        "    self.layer_5 = nn.Linear(400, 300)\n",
        "    self.layer_6 = nn.Linear(300, 1)\n",
        "\n",
        "  def forward(self, x, u):\n",
        "    xu = torch.cat([x, u], 1)\n",
        "    # Forward-Propagation on the first Critic Neural Network\n",
        "    x1 = F.relu(self.layer_1(xu))\n",
        "    x1 = F.relu(self.layer_2(x1))\n",
        "    x1 = self.layer_3(x1)\n",
        "    # Forward-Propagation on the second Critic Neural Network\n",
        "    x2 = F.relu(self.layer_4(xu))\n",
        "    x2 = F.relu(self.layer_5(x2))\n",
        "    x2 = self.layer_6(x2)\n",
        "    return x1, x2\n",
        "\n",
        "  def Q1(self, x, u):\n",
        "    xu = torch.cat([x, u], 1)\n",
        "    x1 = F.relu(self.layer_1(xu))\n",
        "    x1 = F.relu(self.layer_2(x1))\n",
        "    x1 = self.layer_3(x1)\n",
        "    return x1\n",
        "\n",
        "# Selecting the device (CPU or GPU)\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# Building the whole Training Process into a class\n",
        "\n",
        "class TD3(object):\n",
        "  \n",
        "  def __init__(self, state_dim, action_dim, max_action):\n",
        "    self.actor = Actor(state_dim, action_dim, max_action).to(device)\n",
        "    self.actor_target = Actor(state_dim, action_dim, max_action).to(device)\n",
        "    self.actor_target.load_state_dict(self.actor.state_dict())\n",
        "    self.actor_optimizer = torch.optim.Adam(self.actor.parameters())\n",
        "    self.critic = Critic(state_dim, action_dim).to(device)\n",
        "    self.critic_target = Critic(state_dim, action_dim).to(device)\n",
        "    self.critic_target.load_state_dict(self.critic.state_dict())\n",
        "    self.critic_optimizer = torch.optim.Adam(self.critic.parameters())\n",
        "    self.max_action = max_action\n",
        "\n",
        "  def select_action(self, state):\n",
        "    state = torch.Tensor(state.reshape(1, -1)).to(device)\n",
        "    return self.actor(state).cpu().data.numpy().flatten()\n",
        "\n",
        "  def train(self, replay_buffer, iterations, batch_size=100, discount=0.99, tau=0.005, policy_noise=0.2, noise_clip=0.5, policy_freq=2):\n",
        "    \n",
        "    for it in range(iterations):\n",
        "      \n",
        "      # Step 4: We sample a batch of transitions (s, s’, a, r) from the memory\n",
        "      batch_states, batch_next_states, batch_actions, batch_rewards, batch_dones = replay_buffer.sample(batch_size)\n",
        "      state = torch.Tensor(batch_states).to(device)\n",
        "      next_state = torch.Tensor(batch_next_states).to(device)\n",
        "      action = torch.Tensor(batch_actions).to(device)\n",
        "      reward = torch.Tensor(batch_rewards).to(device)\n",
        "      done = torch.Tensor(batch_dones).to(device)\n",
        "      \n",
        "      # Step 5: From the next state s’, the Actor target plays the next action a’\n",
        "      next_action = self.actor_target(next_state)\n",
        "      \n",
        "      # Step 6: We add Gaussian noise to this next action a’ and we clamp it in a range of values supported by the environment\n",
        "      noise = torch.Tensor(batch_actions).data.normal_(0, policy_noise).to(device)\n",
        "      noise = noise.clamp(-noise_clip, noise_clip)\n",
        "      next_action = (next_action + noise).clamp(-self.max_action, self.max_action)\n",
        "      \n",
        "      # Step 7: The two Critic targets take each the couple (s’, a’) as input and return two Q-values Qt1(s’,a’) and Qt2(s’,a’) as outputs\n",
        "      target_Q1, target_Q2 = self.critic_target(next_state, next_action)\n",
        "      \n",
        "      # Step 8: We keep the minimum of these two Q-values: min(Qt1, Qt2)\n",
        "      target_Q = torch.min(target_Q1, target_Q2)\n",
        "      \n",
        "      # Step 9: We get the final target of the two Critic models, which is: Qt = r + γ * min(Qt1, Qt2), where γ is the discount factor\n",
        "      target_Q = reward + ((1 - done) * discount * target_Q).detach()\n",
        "      \n",
        "      # Step 10: The two Critic models take each the couple (s, a) as input and return two Q-values Q1(s,a) and Q2(s,a) as outputs\n",
        "      current_Q1, current_Q2 = self.critic(state, action)\n",
        "      \n",
        "      # Step 11: We compute the loss coming from the two Critic models: Critic Loss = MSE_Loss(Q1(s,a), Qt) + MSE_Loss(Q2(s,a), Qt)\n",
        "      critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)\n",
        "      \n",
        "      # Step 12: We backpropagate this Critic loss and update the parameters of the two Critic models with a SGD optimizer\n",
        "      self.critic_optimizer.zero_grad()\n",
        "      critic_loss.backward()\n",
        "      self.critic_optimizer.step()\n",
        "      \n",
        "      # Step 13: Once every two iterations, we update our Actor model by performing gradient ascent on the output of the first Critic model\n",
        "      if it % policy_freq == 0:\n",
        "        actor_loss = -self.critic.Q1(state, self.actor(state)).mean()\n",
        "        self.actor_optimizer.zero_grad()\n",
        "        actor_loss.backward()\n",
        "        self.actor_optimizer.step()\n",
        "        \n",
        "        # Step 14: Still once every two iterations, we update the weights of the Actor target by polyak averaging\n",
        "        for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):\n",
        "          target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)\n",
        "        \n",
        "        # Step 15: Still once every two iterations, we update the weights of the Critic target by polyak averaging\n",
        "        for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):\n",
        "          target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)\n",
        "  \n",
        "  # Making a save method to save a trained model\n",
        "  def save(self, filename, directory):\n",
        "    torch.save(self.actor.state_dict(), '%s/%s_actor.pth' % (directory, filename))\n",
        "    torch.save(self.critic.state_dict(), '%s/%s_critic.pth' % (directory, filename))\n",
        "  \n",
        "  # Making a load method to load a pre-trained model\n",
        "  def load(self, filename, directory):\n",
        "    self.actor.load_state_dict(torch.load('%s/%s_actor.pth' % (directory, filename)))\n",
        "    self.critic.load_state_dict(torch.load('%s/%s_critic.pth' % (directory, filename)))\n",
        "\n",
        "def evaluate_policy(policy, eval_episodes=10):\n",
        "  avg_reward = 0.\n",
        "  for _ in range(eval_episodes):\n",
        "    obs = env.reset()\n",
        "    done = False\n",
        "    while not done:\n",
        "      action = policy.select_action(np.array(obs))\n",
        "      obs, reward, done, _ = env.step(action)\n",
        "      avg_reward += reward\n",
        "  avg_reward /= eval_episodes\n",
        "  print (\"---------------------------------------\")\n",
        "  print (\"Average Reward over the Evaluation Step: %f\" % (avg_reward))\n",
        "  print (\"---------------------------------------\")\n",
        "  return avg_reward\n",
        "\n",
        "env_name = \"Walker2DBulletEnv-v0\"\n",
        "seed = 0\n",
        "\n",
        "file_name = \"%s_%s_%s\" % (\"TD3\", env_name, str(seed))\n",
        "print (\"---------------------------------------\")\n",
        "print (\"Settings: %s\" % (file_name))\n",
        "print (\"---------------------------------------\")\n",
        "\n",
        "eval_episodes = 10\n",
        "save_env_vid = True\n",
        "env = gym.make(env_name)\n",
        "max_episode_steps = env._max_episode_steps\n",
        "if save_env_vid:\n",
        "  env = wrappers.Monitor(env, monitor_dir, force = True)\n",
        "  env.reset()\n",
        "env.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "np.random.seed(seed)\n",
        "state_dim = env.observation_space.shape[0]\n",
        "action_dim = env.action_space.shape[0]\n",
        "max_action = float(env.action_space.high[0])\n",
        "policy = TD3(state_dim, action_dim, max_action)\n",
        "policy.load(file_name, './pytorch_models/')\n",
        "_ = evaluate_policy(policy, eval_episodes=eval_episodes)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}