TD3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Twin Delayed Deep Deterministic Policy Gradient (TD3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "import torch.nn.functional as F\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "import gym\n",
    "import roboschool\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hidden_init(layer):\n",
    "    fan_in = layer.weight.data.size()[0]\n",
    "    lim = 1. / np.sqrt(fan_in)\n",
    "    return (-lim, lim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Actor(nn.Module):\n",
    "    \"\"\"Initialize parameters and build model.\n",
    "        Args:\n",
    "            state_size (int): Dimension of each state\n",
    "            action_size (int): Dimension of each action\n",
    "            max_action (float): highest action to take\n",
    "            seed (int): Random seed\n",
    "            h1_units (int): Number of nodes in first hidden layer\n",
    "            h2_units (int): Number of nodes in second hidden layer\n",
    "            \n",
    "        Return:\n",
    "            action output of network with tanh activation\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, state_dim, action_dim, max_action):\n",
    "        super(Actor, self).__init__()\n",
    "\n",
    "        self.l1 = nn.Linear(state_dim, 400)\n",
    "        self.l2 = nn.Linear(400, 300)\n",
    "        self.l3 = nn.Linear(300, action_dim)\n",
    "\n",
    "        self.max_action = max_action\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        print(x)\n",
    "        x = F.relu(self.l1(x))\n",
    "        x = F.relu(self.l2(x))\n",
    "        x = self.max_action * torch.tanh(self.l3(x)) \n",
    "        return x\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Critic(nn.Module):\n",
    "    \"\"\"Initialize parameters and build model.\n",
    "        Args:\n",
    "            state_size (int): Dimension of each state\n",
    "            action_size (int): Dimension of each action\n",
    "            max_action (float): highest action to take\n",
    "            seed (int): Random seed\n",
    "            h1_units (int): Number of nodes in first hidden layer\n",
    "            h2_units (int): Number of nodes in second hidden layer\n",
    "            \n",
    "        Return:\n",
    "            value output of network \n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, state_dim, action_dim):\n",
    "        super(Critic, self).__init__()\n",
    "\n",
    "        # Q1 architecture\n",
    "        self.l1 = nn.Linear(state_dim + action_dim, 400)\n",
    "        self.l2 = nn.Linear(400, 300)\n",
    "        self.l3 = nn.Linear(300, 1)\n",
    "\n",
    "        # Q2 architecture\n",
    "        self.l4 = nn.Linear(state_dim + action_dim, 400)\n",
    "        self.l5 = nn.Linear(400, 300)\n",
    "        self.l6 = nn.Linear(300, 1)\n",
    "\n",
    "\n",
    "    def forward(self, x, u):\n",
    "        xu = torch.cat([x, u], 1)\n",
    "\n",
    "        x1 = F.relu(self.l1(xu))\n",
    "        x1 = F.relu(self.l2(x1))\n",
    "        x1 = self.l3(x1)\n",
    "\n",
    "        x2 = F.relu(self.l4(xu))\n",
    "        x2 = F.relu(self.l5(x2))\n",
    "        x2 = self.l6(x2)\n",
    "        return x1, x2\n",
    "\n",
    "\n",
    "    def Q1(self, x, u):\n",
    "        xu = torch.cat([x, u], 1)\n",
    "\n",
    "        x1 = F.relu(self.l1(xu))\n",
    "        x1 = F.relu(self.l2(x1))\n",
    "        x1 = self.l3(x1)\n",
    "        return x1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code based on: \n",
    "# https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py\n",
    "\n",
    "# Expects tuples of (state, next_state, action, reward, done)\n",
    "class ReplayBuffer(object):\n",
    "    \"\"\"Buffer to store tuples of experience replay\"\"\"\n",
    "    \n",
    "    def __init__(self, max_size=1000000):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            max_size (int): total amount of tuples to store\n",
    "        \"\"\"\n",
    "        \n",
    "        self.storage = []\n",
    "        self.max_size = max_size\n",
    "        self.ptr = 0\n",
    "\n",
    "    def add(self, data):\n",
    "        \"\"\"Add experience tuples to buffer\n",
    "        \n",
    "        Args:\n",
    "            data (tuple): experience replay tuple\n",
    "        \"\"\"\n",
    "        \n",
    "        if len(self.storage) == self.max_size:\n",
    "            self.storage[int(self.ptr)] = data\n",
    "            self.ptr = (self.ptr + 1) % self.max_size\n",
    "        else:\n",
    "            self.storage.append(data)\n",
    "\n",
    "    def sample(self, batch_size):\n",
    "        \"\"\"Samples a random amount of experiences from buffer of batch size\n",
    "        \n",
    "        Args:\n",
    "            batch_size (int): size of sample\n",
    "        \"\"\"\n",
    "        \n",
    "        ind = np.random.randint(0, len(self.storage), size=batch_size)\n",
    "        states, actions, next_states, rewards, dones = [], [], [], [], []\n",
    "\n",
    "        for i in ind: \n",
    "            s, a, s_, r, d = self.storage[i]\n",
    "            states.append(np.array(s, copy=False))\n",
    "            actions.append(np.array(a, copy=False))\n",
    "            next_states.append(np.array(s_, copy=False))\n",
    "            rewards.append(np.array(r, copy=False))\n",
    "            dones.append(np.array(d, copy=False))\n",
    "\n",
    "        return np.array(states), np.array(actions), np.array(next_states), np.array(rewards).reshape(-1, 1), np.array(dones).reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TD3(object):\n",
    "    \"\"\"Agent class that handles the training of the networks and provides outputs as actions\n",
    "    \n",
    "        Args:\n",
    "            state_dim (int): state size\n",
    "            action_dim (int): action size\n",
    "            max_action (float): highest action to take\n",
    "            device (device): cuda or cpu to process tensors\n",
    "            env (env): gym environment to use\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, state_dim, action_dim, max_action, env):\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(), lr=1e-3)\n",
    "\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(), lr=1e-3)\n",
    "\n",
    "        self.max_action = max_action\n",
    "        self.env = env\n",
    "\n",
    "\n",
    "        \n",
    "    def select_action(self, state, noise=0.1):\n",
    "        \"\"\"Select an appropriate action from the agent policy\n",
    "        \n",
    "            Args:\n",
    "                state (array): current state of environment\n",
    "                noise (float): how much noise to add to acitons\n",
    "                \n",
    "            Returns:\n",
    "                action (float): action clipped within action range\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        state = torch.FloatTensor(state.reshape(1, -1)).to(device)\n",
    "        print(state)\n",
    "        \n",
    "        action = self.actor(state).cpu().data.numpy().flatten()\n",
    "        if noise != 0: \n",
    "            action = (action + np.random.normal(0, noise, size=self.env.action_space.shape[0]))\n",
    "            \n",
    "        return action.clip(self.env.action_space.low, self.env.action_space.high)\n",
    "\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",
    "        \"\"\"Train and update actor and critic networks\n",
    "        \n",
    "            Args:\n",
    "                replay_buffer (ReplayBuffer): buffer for experience replay\n",
    "                iterations (int): how many times to run training\n",
    "                batch_size(int): batch size to sample from replay buffer\n",
    "                discount (float): discount factor\n",
    "                tau (float): soft update for main networks to target networks\n",
    "                \n",
    "            Return:\n",
    "                actor_loss (float): loss from actor network\n",
    "                critic_loss (float): loss from critic network\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "        for it in range(iterations):\n",
    "\n",
    "            # Sample replay buffer \n",
    "            x, y, u, r, d = replay_buffer.sample(batch_size)\n",
    "            state = torch.FloatTensor(x).to(device)\n",
    "            action = torch.FloatTensor(u).to(device)\n",
    "            next_state = torch.FloatTensor(y).to(device)\n",
    "            done = torch.FloatTensor(1 - d).to(device)\n",
    "            reward = torch.FloatTensor(r).to(device)\n",
    "\n",
    "            # Select action according to policy and add clipped noise \n",
    "            noise = torch.FloatTensor(u).data.normal_(0, policy_noise).to(device)\n",
    "            noise = noise.clamp(-noise_clip, noise_clip)\n",
    "            next_action = (self.actor_target(next_state) + noise).clamp(-self.max_action, self.max_action)\n",
    "\n",
    "            # Compute the target Q value\n",
    "            target_Q1, target_Q2 = self.critic_target(next_state, next_action)\n",
    "            target_Q = torch.min(target_Q1, target_Q2)\n",
    "            target_Q = reward + (done * discount * target_Q).detach()\n",
    "\n",
    "            # Get current Q estimates\n",
    "            current_Q1, current_Q2 = self.critic(state, action)\n",
    "\n",
    "            # Compute critic loss\n",
    "            critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q) \n",
    "\n",
    "            # Optimize the critic\n",
    "            self.critic_optimizer.zero_grad()\n",
    "            critic_loss.backward()\n",
    "            self.critic_optimizer.step()\n",
    "\n",
    "            # Delayed policy updates\n",
    "            if it % policy_freq == 0:\n",
    "\n",
    "                # Compute actor loss\n",
    "                actor_loss = -self.critic.Q1(state, self.actor(state)).mean()\n",
    "\n",
    "                # Optimize the actor \n",
    "                self.actor_optimizer.zero_grad()\n",
    "                actor_loss.backward()\n",
    "                self.actor_optimizer.step()\n",
    "\n",
    "                # Update the frozen target models\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",
    "                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",
    "\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",
    "\n",
    "    def load(self, filename=\"best_avg\", directory=\"./saves\"):\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": [
    "# Runner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Runner():\n",
    "    \"\"\"Carries out the environment steps and adds experiences to memory\"\"\"\n",
    "    \n",
    "    def __init__(self, env, agent, replay_buffer):\n",
    "        \n",
    "        self.env = env\n",
    "        self.agent = agent\n",
    "        self.replay_buffer = replay_buffer\n",
    "        self.obs = env.reset()\n",
    "        self.done = False\n",
    "        \n",
    "    def next_step(self, episode_timesteps, noise=0.1):\n",
    "        \n",
    "        action = self.agent.select_action(np.array(self.obs), noise=0.1)\n",
    "        \n",
    "        # Perform action\n",
    "        new_obs, reward, done, _ = self.env.step(action) \n",
    "        done_bool = 0 if episode_timesteps + 1 == 200 else float(done)\n",
    "    \n",
    "        # Store data in replay buffer\n",
    "        replay_buffer.add((self.obs, new_obs, action, reward, done_bool))\n",
    "        \n",
    "        self.obs = new_obs\n",
    "        \n",
    "        if done:\n",
    "            self.obs = self.env.reset()\n",
    "            done = False\n",
    "            \n",
    "            return reward, True\n",
    "        \n",
    "        return reward, done"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_policy(policy, env, eval_episodes=100,render=False):\n",
    "    \"\"\"run several episodes using the best agent policy\n",
    "        \n",
    "        Args:\n",
    "            policy (agent): agent to evaluate\n",
    "            env (env): gym environment\n",
    "            eval_episodes (int): how many test episodes to run\n",
    "            render (bool): show training\n",
    "        \n",
    "        Returns:\n",
    "            avg_reward (float): average reward over the number of evaluations\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    avg_reward = 0.\n",
    "    for i in range(eval_episodes):\n",
    "        obs = env.reset()\n",
    "        done = False\n",
    "        while not done:\n",
    "            if render:\n",
    "                env.render()\n",
    "            action = policy.select_action(np.array(obs), noise=0)\n",
    "            obs, reward, done, _ = env.step(action)\n",
    "            avg_reward += reward\n",
    "\n",
    "    avg_reward /= eval_episodes\n",
    "\n",
    "    print(\"\\n---------------------------------------\")\n",
    "    print(\"Evaluation over {:d} episodes: {:f}\" .format(eval_episodes, avg_reward))\n",
    "    print(\"---------------------------------------\")\n",
    "    return avg_reward"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Observation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def observe(env,replay_buffer, observation_steps):\n",
    "    \"\"\"run episodes while taking random actions and filling replay_buffer\n",
    "    \n",
    "        Args:\n",
    "            env (env): gym environment\n",
    "            replay_buffer(ReplayBuffer): buffer to store experience replay\n",
    "            observation_steps (int): how many steps to observe for\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    time_steps = 0\n",
    "    obs = env.reset()\n",
    "    done = False\n",
    "\n",
    "    while time_steps < observation_steps:\n",
    "        action = env.action_space.sample()\n",
    "        new_obs, reward, done, _ = env.step(action)\n",
    "\n",
    "        replay_buffer.add((obs, new_obs, action, reward, done))\n",
    "\n",
    "        obs = new_obs\n",
    "        time_steps += 1\n",
    "\n",
    "        if done:\n",
    "            obs = env.reset()\n",
    "            done = False\n",
    "\n",
    "        print(\"\\rPopulating Buffer {}/{}.\".format(time_steps, observation_steps), end=\"\")\n",
    "        sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(agent, test_env):\n",
    "    \"\"\"Train the agent for exploration steps\n",
    "    \n",
    "        Args:\n",
    "            agent (Agent): agent to use\n",
    "            env (environment): gym environment\n",
    "            writer (SummaryWriter): tensorboard writer\n",
    "            exploration (int): how many training steps to run\n",
    "    \n",
    "    \"\"\"\n",
    "\n",
    "    total_timesteps = 0\n",
    "    timesteps_since_eval = 0\n",
    "    episode_num = 0\n",
    "    episode_reward = 0\n",
    "    episode_timesteps = 0\n",
    "    done = False \n",
    "    obs = env.reset()\n",
    "    evaluations = []\n",
    "    rewards = []\n",
    "    best_avg = -2000\n",
    "    \n",
    "    writer = SummaryWriter(comment=\"-TD3_Baseline_HalfCheetah\")\n",
    "    \n",
    "    while total_timesteps < EXPLORATION:\n",
    "    \n",
    "        if done: \n",
    "\n",
    "            if total_timesteps != 0: \n",
    "                rewards.append(episode_reward)\n",
    "                avg_reward = np.mean(rewards[-100:])\n",
    "                \n",
    "                writer.add_scalar(\"avg_reward\", avg_reward, total_timesteps)\n",
    "                writer.add_scalar(\"reward_step\", reward, total_timesteps)\n",
    "                writer.add_scalar(\"episode_reward\", episode_reward, total_timesteps)\n",
    "                \n",
    "                if best_avg < avg_reward:\n",
    "                    best_avg = avg_reward\n",
    "                    print(\"saving best model....\\n\")\n",
    "                    agent.save(\"best_avg\",\"saves\")\n",
    "\n",
    "                print(\"\\rTotal T: {:d} Episode Num: {:d} Reward: {:f} Avg Reward: {:f}\".format(\n",
    "                    total_timesteps, episode_num, episode_reward, avg_reward), end=\"\")\n",
    "                sys.stdout.flush()\n",
    "\n",
    "\n",
    "                if avg_reward >= REWARD_THRESH:\n",
    "                    break\n",
    "\n",
    "                agent.train(replay_buffer, episode_timesteps, BATCH_SIZE, GAMMA, TAU, NOISE, NOISE_CLIP, POLICY_FREQUENCY)\n",
    "\n",
    "                # Evaluate episode\n",
    "                if timesteps_since_eval >= EVAL_FREQUENCY:\n",
    "                    timesteps_since_eval %= EVAL_FREQUENCY\n",
    "                    eval_reward = evaluate_policy(agent, test_env)\n",
    "                    evaluations.append(avg_reward)\n",
    "                    writer.add_scalar(\"eval_reward\", eval_reward, total_timesteps)\n",
    "\n",
    "                    if best_avg < eval_reward:\n",
    "                        best_avg = eval_reward\n",
    "                        print(\"saving best model....\\n\")\n",
    "                        agent.save(\"best_avg\",\"saves\")\n",
    "\n",
    "                episode_reward = 0\n",
    "                episode_timesteps = 0\n",
    "                episode_num += 1 \n",
    "\n",
    "        reward, done = runner.next_step(episode_timesteps)\n",
    "        episode_reward += reward\n",
    "\n",
    "        episode_timesteps += 1\n",
    "        total_timesteps += 1\n",
    "        timesteps_since_eval += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENV = \"Pendulum-v0\"#\"RoboschoolHalfCheetah-v1\"\n",
    "SEED = 0\n",
    "OBSERVATION = 10000\n",
    "EXPLORATION = 5000000\n",
    "BATCH_SIZE = 100\n",
    "GAMMA = 0.99\n",
    "TAU = 0.005\n",
    "NOISE = 0.2\n",
    "NOISE_CLIP = 0.5\n",
    "EXPLORE_NOISE = 0.1\n",
    "POLICY_FREQUENCY = 2\n",
    "EVAL_FREQUENCY = 5000\n",
    "REWARD_THRESH = 8000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(ENV)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Set seeds\n",
    "env.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "\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",
    "\n",
    "policy = TD3(state_dim, action_dim, max_action, env)\n",
    "\n",
    "replay_buffer = ReplayBuffer()\n",
    "\n",
    "runner = Runner(env, policy, replay_buffer)\n",
    "\n",
    "total_timesteps = 0\n",
    "timesteps_since_eval = 0\n",
    "episode_num = 0\n",
    "done = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Populate replay buffer\n",
    "observe(env, replay_buffer, OBSERVATION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train agent\n",
    "train(policy, env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load trained policy\n",
    "policy.load()\n",
    "\n",
    "# watch the trained agent run \n",
    "for i in range(10):\n",
    "    evaluate_policy(policy, env, render=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.close()"
   ]
  },
  {
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}