TD3_Quadruped.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installing the packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pybullet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing the libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import random\n",
    "import numpy as np\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Initialze the Experience Replay Memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the first episode we will populate the Experience Replay Memory with random actions.  \n",
    "Then in each subsequent training episode we will leverage some instances of the Experience Replay Memory to account for anomolous (state, action) pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer(object):\n",
    "    def __init__(self, max_size=1e6):\n",
    "        self.storage = []\n",
    "        self.max_size = max_size\n",
    "        self.ptr = 0\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, \\\n",
    "        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(action, copy=False))\n",
    "            batch_rewards.append(np.array(done, copy=False))\n",
    "            batch_dones.append(np.array(done, copy=False))\n",
    "        return np.array(batch_states), np.array(batch_next_states),\\\n",
    "        np.array(batch_actions), np.array(batch_rewards).reshape(-1,1),\\\n",
    "        np.array(batch_dones).reshape(-1,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Build Neural Network for Actor Model & Actor Target"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An Actor Class is created for easy instantiation during Training.\n",
    "To ensure we have a continuous action-space, as opposed to a discrete action space, we map the result of the *tanh() function* to the maximum allowed action-value, ie;  **self.max_action**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Actor(nn.Module):\n",
    "    def __init__(self, state_dim, action_dim, max_action):\n",
    "        super(Actor, self).__init__() # To activate inheritence\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): ##Forward Propagation\n",
    "        x = F.relu(self.layer_1(x))    # F = Functional Module\n",
    "        x = F.relu(self.layer_2(x))\n",
    "        x = self.max_action *torch.tanh(self.layer_3(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Build Neural Network for Critic Models & Critic Targets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We add the second forward method to obtain the gradient ascent forward propagation of the first critic neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "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",
    "\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(x))\n",
    "        x1 = self.layer_3(x)\n",
    "        return x1\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Steps 4 to 15: Training Process"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We run a full episode with first 10,000 actions played randomly to facilitate exploration. Then play actions by the Actor model. Only after that do we start to sample from the Experience Replay Memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "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)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We make a function that evaluates the policy by calculating its average reward over 10 episodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We set the parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_name = \"AntBulletEnv-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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We create a file name for the two saved models: the Actor and Critic models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------------------------\n",
      "Settings: TD3_AntBulletEnv-v0_0\n",
      "---------------------------------------\n"
     ]
    }
   ],
   "source": [
    "file_name = \"%s_%s_%s\" % (\"TD3\", env_name, str(seed))\n",
    "print (\"---------------------------------------\")\n",
    "print (\"Settings: %s\" % (file_name))\n",
    "print (\"---------------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We create a folder inside which we savethe trained models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "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\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We create the PyBullet environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(env_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We set seeds and we get the necessary information on the states and actions in the chosen environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "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])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We create the policy network (the Actor model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = TD3(state_dim, action_dim, max_action)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We create the Experience Replay Memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "replay_buffer = ReplayBuffer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We define a list where all the evaluation results over 10 episodes are stored"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------\n",
      "Average Reward over the Evaluation Step: 9.804960\n",
      "-------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "evaluations = [evaluate_policy(policy)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Folder to store the final results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We initialize the variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_timesteps = 0\n",
    "timesteps_since_eval = 0\n",
    "episode_num = 0\n",
    "done = True\n",
    "t0 = time.time()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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 = \"AntBulletEnv-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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}