6_twin_delayed_ddpg.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "StmTP2EXwnUr"
      },
      "source": [
        "## Twin Delayed DDPG (TD3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZOSJl-X7zvs4"
      },
      "source": [
        "#### Setup virtual display"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "B-Z6takfzqGk"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<pyvirtualdisplay.display.Display at 0x7f9ef4c60eb0>"
            ]
          },
          "execution_count": 1,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "from pyvirtualdisplay import Display\n",
        "Display(visible=False, size=(1400, 900)).start()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cz8DLleGz_TF"
      },
      "source": [
        "#### Import the necessary code libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "cP5t6U7-nYoc"
      },
      "outputs": [],
      "source": [
        "import copy\n",
        "import gym\n",
        "import torch\n",
        "import itertools\n",
        "\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from collections import deque, namedtuple\n",
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "\n",
        "from torch import nn\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.data.dataset import IterableDataset\n",
        "from torch.optim import AdamW\n",
        "\n",
        "from pytorch_lightning import LightningModule, Trainer\n",
        "\n",
        "from gym.wrappers import RecordVideo, RecordEpisodeStatistics\n",
        "\n",
        "\n",
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "num_gpus = torch.cuda.device_count()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "import copy\n",
        "import gym\n",
        "import torch\n",
        "import random\n",
        "import functools\n",
        "import itertools\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from collections import deque, namedtuple\n",
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "\n",
        "from torch import nn\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.data.dataset import IterableDataset\n",
        "from torch.optim import AdamW\n",
        "\n",
        "from pytorch_lightning import LightningModule, Trainer\n",
        "\n",
        "import brax\n",
        "from brax.v1 import envs\n",
        "from brax.v1.envs.to_torch import JaxToTorchWrapper\n",
        "\n",
        "# from brax.envs import to_torch#this will ensure that our environments can use pytorch tensors\n",
        "from brax.v1.io import html#this will allow us to display our environments in the notebook\n",
        "# from gym.wrappers import RecordVideo, RecordEpisodeStatistics\n",
        "from brax.v1.io import image\n",
        "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "num_gpus = torch.cuda.device_count()\n",
        "\n",
        "v = torch.ones(1, device='cuda')#if we plan on using the gpu, we need to run this line, \n",
        "# #we have to create this tensor before we do anything with brax otherwise, brax will be a little bit greedy and it will suck up memory from gpu\n",
        "import jax\n",
        "\n",
        "jax.config.update('jax_platform_name', 'gpu')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "Z_IrPlU1wwPx"
      },
      "outputs": [],
      "source": [
        "def display_video(episode=0):\n",
        "  video_file = open(f'/content/videos/rl-video-episode-{episode}.mp4', \"r+b\").read()\n",
        "  video_url = f\"data:video/mp4;base64,{b64encode(video_file).decode()}\"\n",
        "  return HTML(f\"<video width=600 controls><source src='{video_url}'></video>\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dj6w5kLkz9DN"
      },
      "source": [
        "#### Create the experience buffer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "E18jarh3ybBW"
      },
      "outputs": [],
      "source": [
        "Experience = namedtuple(\n",
        "    \"Experience\",\n",
        "    field_names=[\"state\", \"action\", \"reward\", \"done\", \"new_state\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fw-77TRyBvHz"
      },
      "outputs": [],
      "source": [
        "class ReplayBuffer:\n",
        "\n",
        "  def __init__(self, capacity):\n",
        "      self.buffer = deque(maxlen=capacity)\n",
        "\n",
        "  def __len__(self):\n",
        "      return len(self.buffer)\n",
        "\n",
        "  def append(self, experience):\n",
        "\n",
        "      self.buffer.append(experience)\n",
        "\n",
        "  def sample(self, batch_size):\n",
        "      indices = np.random.choice(len(self.buffer), batch_size, replace=False)\n",
        "      states, actions, rewards, dones, next_states = zip(*(self.buffer[idx] for idx in indices))\n",
        "\n",
        "      return (\n",
        "          np.array(states, dtype=np.float32),\n",
        "          np.array(actions),\n",
        "          np.array(rewards, dtype=np.float32),\n",
        "          np.array(dones, dtype=np.bool),\n",
        "          np.array(next_states, dtype=np.float32)\n",
        "      )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [],
      "source": [
        "class ReplayBuffer:\n",
        "\n",
        "  def __init__(self, capacity):\n",
        "    self.buffer = deque(maxlen=capacity)\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.buffer)\n",
        "  \n",
        "  def append(self, experience):\n",
        "    self.buffer.append(experience)\n",
        "  \n",
        "  def sample(self, batch_size):\n",
        "    return random.sample(self.buffer, batch_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YtGko6LVQaJz"
      },
      "outputs": [],
      "source": [
        "class RLDataset(IterableDataset):\n",
        "\n",
        "  def __init__(self, buffer, sample_size=200):\n",
        "      self.buffer = buffer\n",
        "      self.sample_size = sample_size\n",
        "\n",
        "  def __iter__(self):\n",
        "      states, actions, rewards, dones, new_states = self.buffer.sample(self.sample_size)\n",
        "      for i in range(len(dones)):\n",
        "          yield states[i], actions[i], rewards[i], dones[i], new_states[i]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {},
      "outputs": [],
      "source": [
        "class RLDataset(IterableDataset):\n",
        "\n",
        "  def __init__(self, buffer, sample_size=400):\n",
        "    self.buffer = buffer\n",
        "    self.sample_size = sample_size\n",
        "  \n",
        "  def __iter__(self):\n",
        "    for experience in self.buffer.sample(self.sample_size):\n",
        "      yield experience"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ihkyoL5WQgGm"
      },
      "source": [
        "#### Create the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "entry_point = functools.partial(envs.create_gym_env, env_name='ant')#we are doing this to interface this with the gym library so that \n",
        "#we don't have to change our coding style that we have learnt with the gym library\n",
        "gym.register('brax-ant-v0', entry_point=entry_point)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "def create_environment(env_name, num_envs=256, episode_length=1000):\n",
        "  env = gym.make(env_name, batch_size=num_envs, episode_length=episode_length)#in physics based simulation engines, we can pass a parameter called batch_size which will decide \n",
        "  #how many copies of the environment will be running at the same time. This will help us create the required replay buffer and the episode_length is like a limit on the number of steps\n",
        "  #or more simply, how long we wish for an episode to run\n",
        "  # env = RecordVideo(env, video_folder='./videos2', episode_trigger=lambda x: x % 100 == 0)\n",
        "  env = JaxToTorchWrapper(env, device=device)#brax is a physics engine tool that uses a different numerical computing tool to run simulation called JAX but our tool of choice is called \n",
        "  #pytorch, so by callling this wrapper, the environment will take up our pytorch tensors, convert them into the right format, use them internally, and then speed back pytorch sensors so that we can work with that\n",
        "  return env"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [],
      "source": [
        "@torch.no_grad()\n",
        "def test_env(env_name, policy=None):\n",
        "  env = gym.make(env_name, episode_length=1000)\n",
        "  env = JaxToTorchWrapper(env, device=device)\n",
        "  # print(env)\n",
        "  # print(env.unwrapped.unwrapped)\n",
        "  # print(env.unwrapped._state)\n",
        "  # print(env.unwrapped._env)\n",
        "  qp_array = []\n",
        "  state = env.reset()\n",
        "  for i in range(1000):\n",
        "    if policy:\n",
        "      action = policy.net(state.unsqueeze(0)).squeeze()\n",
        "    else:\n",
        "      action = env.action_space.sample()\n",
        "    state, _, _, _ = env.step(action)\n",
        "    qp_array.append(env.unwrapped._state.qp)\n",
        "  return HTML(html.render(env.unwrapped._env.sys, qp_array))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xZ3h8CCOQjGL"
      },
      "outputs": [],
      "source": [
        "def create_environment(name):\n",
        "  env = gym.make(name)\n",
        "  env = RecordVideo(env, video_folder='./videos', episode_trigger=lambda x: x % 50 == 0)\n",
        "  env = RecordEpisodeStatistics(env)\n",
        "  return env"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b8fKGgFzQ4EX"
      },
      "source": [
        "#### Update the target network"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "-q-OJaPnBvKf"
      },
      "outputs": [],
      "source": [
        "def polyak_average(net, target_net, tau=0.01):\n",
        "    for qp, tp in zip(net.parameters(), target_net.parameters()):\n",
        "        tp.data.copy_(tau * qp.data + (1 - tau) * tp.data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YQMvxcsly9W0"
      },
      "source": [
        "#### Create the gradient policy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "tk9xqIipf0Bu"
      },
      "outputs": [],
      "source": [
        "class GradientPolicy(nn.Module):\n",
        "\n",
        "  def __init__(self, hidden_size, obs_size, out_dims, min, max):\n",
        "    super().__init__()\n",
        "    self.min = torch.from_numpy(min).to(device)\n",
        "    self.max = torch.from_numpy(max).to(device)\n",
        "    self.net = nn.Sequential(\n",
        "        nn.Linear(obs_size, hidden_size),\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(hidden_size, hidden_size),\n",
        "        nn.ReLU(),           \n",
        "        nn.Linear(hidden_size, out_dims),\n",
        "    )\n",
        "\n",
        "  def mu(self, x):\n",
        "    if isinstance(x, np.ndarray):\n",
        "      x = torch.from_numpy(x).to(device)\n",
        "    return self.net(x)\n",
        "  \n",
        "  def forward(self, x, epsilon=0.0, noise_clip=None):\n",
        "    mu = self.mu(x)\n",
        "    noise = torch.normal(0, epsilon, mu.size(), device=mu.device)\n",
        "    if noise_clip is not None:\n",
        "      noise = torch.clamp(noise, - noise_clip, noise_clip)\n",
        "    mu = mu + noise\n",
        "    action = torch.max(torch.min(mu, self.max), self.min)\n",
        "    action = action.detach().cpu().numpy()\n",
        "    return action\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qZBeQSNvzYRk"
      },
      "source": [
        "Create the Deep Q-Learning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "JxCtS_Pkzb-4"
      },
      "outputs": [],
      "source": [
        "class DQN(nn.Module):\n",
        "\n",
        "  def __init__(self, hidden_size, obs_size, out_dims):\n",
        "    super().__init__()\n",
        "    self.net = nn.Sequential(\n",
        "        nn.Linear(obs_size + out_dims, hidden_size),\n",
        "        nn.ReLU(),\n",
        "        nn.Linear(hidden_size, hidden_size),\n",
        "        nn.ReLU(),           \n",
        "        nn.Linear(hidden_size, 1),\n",
        "    )\n",
        "\n",
        "  def forward(self, state, action):\n",
        "    if isinstance(state, np.ndarray):\n",
        "      state = torch.from_numpy(state).to(device)\n",
        "    if isinstance(action, np.ndarray):\n",
        "      action = torch.from_numpy(action).to(device)\n",
        "    in_vector = torch.hstack((state, action))\n",
        "    return self.net(in_vector)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SKFUd7jAf0Ej"
      },
      "outputs": [],
      "source": [
        "class TD3(LightningModule):\n",
        "\n",
        "  def __init__(self, env_name, capacity=1_000_000, batch_size=256, lr=1e-4, hidden_size=128,\n",
        "               gamma=0.99, loss_fn=F.smooth_l1_loss, optim=AdamW, eps_start=5.0, \n",
        "               eps_end=0.2, eps_last_episode=200, samples_per_epoch=10_000, tau=0.01):\n",
        "\n",
        "    super().__init__()\n",
        "\n",
        "    self.env = create_environment(env_name)\n",
        "\n",
        "    obs_size = self.env.observation_space.shape[0]\n",
        "    action_dims = self.env.action_space.shape[0]\n",
        "    max_action = self.env.action_space.high\n",
        "    min_action = self.env.action_space.low\n",
        "\n",
        "\n",
        "    self.q_net1 = DQN(hidden_size, obs_size, action_dims).to(device)\n",
        "    self.q_net2 = DQN(hidden_size, obs_size, action_dims).to(device)\n",
        "    self.policy = GradientPolicy(hidden_size, obs_size, action_dims, min_action, max_action).to(device)\n",
        "\n",
        "    self.target_policy = copy.deepcopy(self.policy)\n",
        "    self.target_q_net1 = copy.deepcopy(self.q_net1)\n",
        "    self.target_q_net2 = copy.deepcopy(self.q_net2)\n",
        "\n",
        "    self.buffer = ReplayBuffer(capacity=capacity)\n",
        "\n",
        "    self.save_hyperparameters()\n",
        "\n",
        "    while len(self.buffer) < self.hparams.samples_per_epoch:\n",
        "\n",
        "      print(f\"{len(self.buffer)} samples in experience buffer. Filling...\")\n",
        "      \n",
        "      self.play_episodes(epsilon=self.hparams.eps_start)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def play_episodes(self, policy=None, epsilon=0.):\n",
        "      obs = self.env.reset()\n",
        "      done = False\n",
        "\n",
        "      while not done:\n",
        "        if policy:\n",
        "          action = self.policy(obs, epsilon=epsilon)\n",
        "        else:\n",
        "          action = self.env.action_space.sample()\n",
        "          \n",
        "        next_obs, reward, done, info = self.env.step(action)\n",
        "        exp = Experience(obs, action, reward, done, next_obs)\n",
        "        self.buffer.append(exp)\n",
        "        obs = next_obs\n",
        "\n",
        "  def forward(self, x):\n",
        "    output = self.policy(x)\n",
        "    return output\n",
        "\n",
        "  def configure_optimizers(self):\n",
        "    q_net_parameters = itertools.chain(self.q_net1.parameters(), self.q_net2.parameters())\n",
        "    q_net_optimizer = self.hparams.optim(q_net_parameters, lr=self.hparams.lr)\n",
        "    policy_optimizer = self.hparams.optim(self.policy.parameters(), lr=self.hparams.lr)\n",
        "    return [q_net_optimizer, policy_optimizer]\n",
        "\n",
        "  def train_dataloader(self):\n",
        "    dataset = RLDataset(self.buffer, self.hparams.samples_per_epoch)\n",
        "    dataloader = DataLoader(\n",
        "        dataset=dataset,\n",
        "        batch_size=self.hparams.batch_size,\n",
        "    )\n",
        "    return dataloader\n",
        "\n",
        "  def training_step(self, batch, batch_idx, optimizer_idx):\n",
        "    states, actions, rewards, dones, next_states = batch\n",
        "    rewards = rewards.unsqueeze(1)\n",
        "    dones = dones.unsqueeze(1)\n",
        "\n",
        "    if optimizer_idx == 0:\n",
        "\n",
        "      epsilon = max(\n",
        "        self.hparams.eps_end,\n",
        "        self.hparams.eps_start - self.current_epoch / self.hparams.eps_last_episode\n",
        "      )\n",
        "\n",
        "      state_action_values1 = self.q_net1(states, actions)\n",
        "      state_action_values2 = self.q_net2(states, actions)\n",
        "\n",
        "      target_actions = self.target_policy(next_states, epsilon=epsilon, noise_clip=0.5)\n",
        "\n",
        "      next_state_values = torch.min(\n",
        "          self.target_q_net1(next_states, target_actions),\n",
        "          self.target_q_net2(next_states, target_actions)\n",
        "      )\n",
        "\n",
        "      next_state_values[dones] = 0.0\n",
        "\n",
        "      expected_state_action_values = rewards + self.hparams.gamma * next_state_values\n",
        "      q_loss1 = self.hparams.loss_fn(state_action_values1, expected_state_action_values)\n",
        "      q_loss2 = self.hparams.loss_fn(state_action_values2, expected_state_action_values)\n",
        "      total_loss = q_loss1 + q_loss2\n",
        "      self.log(\"episode/MSE Loss\", total_loss)\n",
        "      return total_loss\n",
        "\n",
        "    elif optimizer_idx == 1 and batch_idx % 2 == 0:\n",
        "      policy_loss = - self.q_net1(states, self.policy.mu(states)).mean()\n",
        "      self.log(\"episode/Policy Loss\", policy_loss)\n",
        "      return policy_loss\n",
        "\n",
        "  def training_epoch_end(self, training_step_outputs):\n",
        "    epsilon = max(\n",
        "        self.hparams.eps_end,\n",
        "        self.hparams.eps_start - self.current_epoch / self.hparams.eps_last_episode\n",
        "    )\n",
        "\n",
        "    self.play_episodes(policy=self.policy, epsilon=epsilon)\n",
        "\n",
        "    polyak_average(self.q_net1, self.target_q_net1, tau=self.hparams.tau)\n",
        "    polyak_average(self.q_net2, self.target_q_net2, tau=self.hparams.tau)\n",
        "    polyak_average(self.policy, self.target_policy, tau=self.hparams.tau)\n",
        "\n",
        "    self.log(\"episode/Episode return\", self.env.return_queue[-1], prog_bar=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [],
      "source": [
        "class TD3(LightningModule):\n",
        "\n",
        "  def __init__(self, env_name, capacity=500, batch_size=8192, actor_lr=1e-3, \n",
        "               critic_lr=1e-3, hidden_size=256, gamma=0.99, loss_fn=F.smooth_l1_loss, \n",
        "               optim=AdamW, eps_start=1.0, eps_end=0.2, eps_last_episode=500, \n",
        "               samples_per_epoch=10, tau=0.005):\n",
        "\n",
        "    super().__init__()\n",
        "    self.automatic_optimization = False  # Disable automatic optimization\n",
        "    self.env = create_environment(env_name, num_envs=batch_size)#this time, batch_size denotes the number of paraller environments that will run parallely\n",
        "    #each entry in the the replay buffer will have 8192 entries in itself\n",
        "    # print(type(self.env))\n",
        "    # print(type(self.env.reset()))\n",
        "    self.obs = self.env.reset()\n",
        "    self.videos = []\n",
        "\n",
        "    obs_size = self.env.observation_space.shape[1]#why like this?\n",
        "    #remember, now we are learning parallel environments\n",
        "    #if the observation of a single environment has a size of 8, then \n",
        "    #then shape of the observation space would be [8192, 8]\n",
        "    action_dims = self.env.action_space.shape[1]\n",
        "    #same reason as earlier as to why we are taking the second index this time\n",
        "    max_action = self.env.action_space.high\n",
        "    min_action = self.env.action_space.low\n",
        "\n",
        "    self.q_net_1 = DQN(hidden_size, obs_size, action_dims)\n",
        "    self.q_net_2 = DQN(hidden_size, obs_size, action_dims)\n",
        "    self.policy = GradientPolicy(hidden_size, obs_size, action_dims, min_action, max_action)\n",
        "\n",
        "    self.target_policy = copy.deepcopy(self.policy)\n",
        "    self.target_q_net_1 = copy.deepcopy(self.q_net_1)\n",
        "    self.target_q_net_2 = copy.deepcopy(self.q_net_2)\n",
        "\n",
        "    self.buffer = ReplayBuffer(capacity=capacity)\n",
        "\n",
        "    self.save_hyperparameters()\n",
        "\n",
        "    while len(self.buffer) < self.hparams.samples_per_epoch:\n",
        "      print(f\"{len(self.buffer)} samples in experience buffer. Filling...\")\n",
        "      self.play_episode(epsilon=self.hparams.eps_start)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def play_episode(self, policy=None, epsilon=0.):\n",
        "    #in the previous sections, we have the similar method, where we play an entire episode to sample an experience\n",
        "    #but this time, every time we take an action in the environment, we get 8192 different observations.\n",
        "    #so, instead of playing an episode every time, when we call the play_episode function, we'll simply play one move\n",
        "    if policy:\n",
        "      action = policy(self.obs, epsilon=epsilon)\n",
        "    else:\n",
        "      action = self.env.action_space.sample()\n",
        "    next_obs, reward, done, info = self.env.step(action)\n",
        "    exp = (self.obs, action, reward, done, next_obs)\n",
        "    self.buffer.append(exp)\n",
        "    self.obs = next_obs\n",
        "    return reward.mean()\n",
        "\n",
        "  def forward(self, x):#what happens when someone calls this DDPG class on a state\n",
        "    output = self.policy(x)\n",
        "    return output\n",
        "\n",
        "  def configure_optimizers(self):\n",
        "    #the following line of code was added due to the following error\n",
        "    #RuntimeError: Training with multiple optimizers is only supported with manual optimization. Remove the `optimizer_idx` argument from `training_step`, set `self.automatic_optimization = False` and access your optimizers in `training_step` with `opt1, opt2, ... = self.optimizers()`.\n",
        "    # q_net_params = itertools(self.q_net_1.parameters(),self.q_net_2.parameters())\n",
        "    q_net_1_optimizer = self.hparams.optim(self.q_net_1.parameters(), lr=self.hparams.critic_lr)\n",
        "    q_net_2_optimizer = self.hparams.optim(self.q_net_2.parameters(), lr=self.hparams.critic_lr)\n",
        "    policy_optimizer = self.hparams.optim(self.policy.parameters(), lr=self.hparams.actor_lr)\n",
        "    return [q_net_1_optimizer, q_net_2_optimizer, policy_optimizer]#in this alogrithm, we'll call the training step method, twice, once with the first optimiser and once with the second optimiser\n",
        "    #so that both the actor and the critic have the opportunity to update their weights with the same batch of data\n",
        "\n",
        "  def train_dataloader(self):\n",
        "    dataset = RLDataset(self.buffer, self.hparams.samples_per_epoch)\n",
        "    dataloader = DataLoader(\n",
        "        dataset=dataset,\n",
        "        batch_size=1,#because we know every observation in our replay buffer will have 8192 observations, so each individual observation we store will effectively be a batch\n",
        "        #we make this value 1 to avoid making batches of batches\n",
        "    )\n",
        "    return dataloader\n",
        "\n",
        "  def training_step(self, batch, batch_idx):#since we have an optimiser for the q-network and an optimiser for the policy, the training step will be called twice for every batch of data\n",
        "    #once with the first optimiser, that is, optimiser of the q-network, and another time with the optimiser of the policy network\n",
        "    epsilon = max(\n",
        "        self.hparams.eps_end,\n",
        "        self.hparams.eps_start - self.current_epoch / self.hparams.eps_last_episode\n",
        "    )\n",
        "\n",
        "    mean_reward = self.play_episode(policy=self.policy, epsilon=epsilon)#remember, play episode is now just a single step, but not in one environment, but in 8192 environments\n",
        "    self.log(\"episode/mean_reward\", mean_reward)\n",
        "\n",
        "    \n",
        "\n",
        "    \n",
        "    states, actions, rewards, dones, next_states = map(torch.squeeze, batch)\n",
        "    #what are we doing here? Remember that the dataloader loads batches of data. This time, it will load a batch of 1 because in a single observation, we already have a batch of data\n",
        "    #but that will give us the data in the wrong format. It will give us a batch of data with the size (1, 8192, ..), but it is wrong because 8192 is already the correct size, what we need is (8192, ...)\n",
        "    #so what we will do is apply torch.squeez to each element of the batch using the map function, the squeeze operation will look for indices, where we only have one item and therefore, it's a dimension we can kill\n",
        "    #the problem is now rewards and dones are flat lists\n",
        "    #we'll do the same thing as we did in previous sections\n",
        "    rewards = rewards.unsqueeze(1)\n",
        "    dones = dones.unsqueeze(1).bool()\n",
        "\n",
        "    #the following are changes that are added by me\n",
        "    opt_q_net_1, opt_q_net_2, opt_policy = self.optimizers()\n",
        "    # if isinstance(opt_q_net, torch.optim.Optimizer):\n",
        "    # print('Inside q_network loss optimsation')\n",
        "    next_actions = self.target_policy(next_states, epsilon=epsilon, noise_clip=0.05)\n",
        "    state_action_values_1 = self.q_net_1(states, actions)\n",
        "    state_action_values_2 = self.q_net_2(states, actions)\n",
        "    next_state_values=torch.min(self.target_q_net_1(next_states, next_actions), self.target_q_net_2(next_states, next_actions))\n",
        "    # next_state_values = \n",
        "    next_state_values[dones] = 0.0\n",
        "    expected_state_action_values = rewards + self.hparams.gamma * next_state_values\n",
        "    q_loss_1 = self.hparams.loss_fn(state_action_values_1, expected_state_action_values)\n",
        "    q_loss_2 = self.hparams.loss_fn(state_action_values_2, expected_state_action_values)\n",
        "    opt_q_net_1.zero_grad()\n",
        "    q_loss_1.backward(retain_graph=True)\n",
        "    opt_q_net_1.step()\n",
        "    opt_q_net_2.zero_grad()\n",
        "    q_loss_2.backward()\n",
        "    opt_q_net_2.step()\n",
        "    # self.log_dict({\"episode/Q-Loss\": q_loss})\n",
        "    # return q_loss\n",
        "    \n",
        "      # elif isinstance(opt_policy, torch.optim.Optimizer):\n",
        "      # print(\"Inside policy network loss optimisation\")\n",
        "      # we can pick any of the two networks to optimise the policy\n",
        "    if batch_idx % 2 == 0:#we only want the policy to update half the times the q-networks are updated\n",
        "      mu = self.policy.mu(states)\n",
        "      policy_loss = - self.q_net_1(states, mu).mean()#what exactly are we doing here?\n",
        "      opt_policy.zero_grad()\n",
        "      policy_loss.backward()\n",
        "      opt_policy.step()\n",
        "      #the better the actions of the policy, the higher will be the q-values of the actions taken by that policy in a specific set of stateas\n",
        "      #so we want to increase the values produced by the q-network by modifying only the actions taken by the policy\n",
        "      #that is, here the q-network remains constant, and the better our actions, the higher the estimates of the q-newtork of the values of those actions\n",
        "      #so by passing those actions through the q-network and trying to maximise their value, will be improving the performance of the policy\n",
        "      #but pytorch doesn't have a mechanism to maximise a value, it can only minismise them. Therefore, what we did is to place the negative sign of those values\n",
        "      #because minimising the negative of the value will achieve the same thing as maximising the values themselves. Then we compute the mean of those values and then we have the loss of our policy\n",
        "      self.log_dict({\"episode/Policy Loss\": policy_loss})\n",
        "      # return q_loss, policy_loss\n",
        "  \n",
        "  def on_train_epoch_end(self):\n",
        "    polyak_average(self.q_net_1, self.target_q_net_1, tau=self.hparams.tau)\n",
        "    polyak_average(self.q_net_2, self.target_q_net_2, tau=self.hparams.tau)\n",
        "    polyak_average(self.policy, self.target_policy, tau=self.hparams.tau)\n",
        "    if self.current_epoch % 100 == 0:\n",
        "      video = test_env(self.env.spec.id, policy=self.policy)\n",
        "      self.videos.append(video)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-1b2qnsACnz4"
      },
      "outputs": [],
      "source": [
        "# Start tensorboard.\n",
        "!rm -r /content/lightning_logs/\n",
        "!rm -r /content/videos/\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir /content/lightning_logs/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3rn7R1Sff0HS"
      },
      "outputs": [],
      "source": [
        "algo = TD3('LunarLanderContinuous-v2')\n",
        "\n",
        "trainer = Trainer(\n",
        "    # gpus=num_gpus, \n",
        "    max_epochs=2_000, \n",
        "    track_grad_norm=2,\n",
        ")\n",
        "\n",
        "trainer.fit(algo)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "algo = DDPG('brax-ant-v0')\n",
        "\n",
        "trainer = Trainer(\n",
        "    # gpus=num_gpus, \n",
        "    max_epochs=10_000,\n",
        "    log_every_n_steps=10\n",
        ")\n",
        "\n",
        "trainer.fit(algo)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xd3Qvne4xbHX"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}