deep_rl.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOc6WZqk6AhniJ5oO1y7Csp",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/worthlessFella/deep_learning/blob/main/deep_rl.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "faL-z-KKtxzR"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "class QNetwork(nn.Module):\n",
        "    def __init__(self, state_dim, action_dim, hidden_dim=64):\n",
        "        super(QNetwork, self).__init__()\n",
        "        self.fc1 = nn.Linear(state_dim, hidden_dim)\n",
        "        self.fc2 = nn.Linear(hidden_dim, hidden_dim)\n",
        "        self.fc3 = nn.Linear(hidden_dim, action_dim)\n",
        "\n",
        "    def forward(self, state):\n",
        "        x = torch.relu(self.fc1(state))\n",
        "        x = torch.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x\n",
        "\n",
        "class DQNAgent:\n",
        "    def __init__(self, state_dim, action_dim, gamma=0.99, epsilon=1.0, epsilon_min=0.01, epsilon_decay=0.995, lr=0.001, batch_size=64):\n",
        "        self.state_dim = state_dim\n",
        "        self.action_dim = action_dim\n",
        "        self.gamma = gamma\n",
        "        self.epsilon = epsilon\n",
        "        self.epsilon_min = epsilon_min\n",
        "        self.epsilon_decay = epsilon_decay\n",
        "        self.lr = lr\n",
        "        self.batch_size = batch_size\n",
        "\n",
        "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "        self.q_network = QNetwork(state_dim, action_dim).to(self.device)\n",
        "        self.optimizer = optim.Adam(self.q_network.parameters(), lr=lr)\n",
        "        self.loss_fn = nn.MSELoss()\n",
        "\n",
        "        self.memory = []\n",
        "        self.steps = 0\n",
        "\n",
        "    def act(self, state):\n",
        "        if random.random() < self.epsilon:\n",
        "            return random.randrange(self.action_dim)\n",
        "        state = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(self.device)\n",
        "        with torch.no_grad():\n",
        "            q_values = self.q_network(state)\n",
        "        return torch.argmax(q_values, dim=1).item()\n",
        "\n",
        "    def remember(self, state, action, reward, next_state, done):\n",
        "        self.memory.append((state, action, reward, next_state, done))\n",
        "        self.steps += 1\n",
        "        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)\n",
        "\n",
        "    def replay(self):\n",
        "        if len(self.memory) < self.batch_size:\n",
        "            return\n",
        "        batch = random.sample(self.memory, self.batch_size)\n",
        "        states, actions, rewards, next_states, dones = zip(*batch)\n",
        "        states = torch.tensor(states, dtype=torch.float32).to(self.device)\n",
        "        actions = torch.tensor(actions, dtype=torch.long).unsqueeze(1).to(self.device)\n",
        "        rewards = torch.tensor(rewards, dtype=torch.float32).unsqueeze(1).to(self.device)\n",
        "        next_states = torch.tensor(next_states, dtype=torch.float32).to(self.device)\n",
        "        dones = torch.tensor(dones, dtype=torch.float32).unsqueeze(1).to(self.device)\n",
        "\n",
        "        q_values = self.q_network(states).gather(1, actions)\n",
        "        next_q_values = self.q_network(next_states).max(dim=1, keepdim=True).values.detach()\n",
        "        target_q_values = rewards + (1 - dones) * self.gamma * next_q_values\n",
        "        loss = self.loss_fn(q_values, target_q_values)\n",
        "\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        self.optimizer.step()\n",
        "\n",
        "    def save(self, path):\n",
        "        torch.save(self.q_network.state_dict(), path)\n",
        "\n",
        "    def load(self, path):\n",
        "        self.q_network.load_state_dict(torch.load(path))\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import gym\n",
        "import torch\n",
        "\n",
        "env = gym.make('CartPole-v0')\n",
        "env.seed(0)\n",
        "torch.manual_seed(0)\n",
        "\n",
        "agent = DQNAgent(env.observation_space.shape[0], env.action_space.n)\n",
        "\n",
        "for episode in range(100):\n",
        "    state = env.reset()\n",
        "    total_reward = 0\n",
        "\n",
        "    for step in range(200):\n",
        "        action = agent.act(state)\n",
        "        next_state, reward, done, _ = env.step(action)\n",
        "        total_reward += reward\n",
        "        agent.remember(state, action, reward, next_state, done)\n",
        "        agent.replay()\n",
        "\n",
        "        if done:\n",
        "            break\n",
        "\n",
        "        state = next_state\n",
        "\n",
        "    print(f'Episode {episode}: Total reward = {total_reward}')\n",
        "\n",
        "env.close()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "52yJQLIYuAtR",
        "outputId": "f2794742-5b1a-403a-cdd2-0a214c847c7a"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Episode 0: Total reward = 9.0\n",
            "Episode 1: Total reward = 25.0\n",
            "Episode 2: Total reward = 19.0\n",
            "Episode 3: Total reward = 11.0\n",
            "Episode 4: Total reward = 16.0\n",
            "Episode 5: Total reward = 16.0\n",
            "Episode 6: Total reward = 13.0\n",
            "Episode 7: Total reward = 18.0\n",
            "Episode 8: Total reward = 11.0\n",
            "Episode 9: Total reward = 19.0\n",
            "Episode 10: Total reward = 14.0\n",
            "Episode 11: Total reward = 13.0\n",
            "Episode 12: Total reward = 9.0\n",
            "Episode 13: Total reward = 29.0\n",
            "Episode 14: Total reward = 9.0\n",
            "Episode 15: Total reward = 21.0\n",
            "Episode 16: Total reward = 10.0\n",
            "Episode 17: Total reward = 10.0\n",
            "Episode 18: Total reward = 11.0\n",
            "Episode 19: Total reward = 10.0\n",
            "Episode 20: Total reward = 11.0\n",
            "Episode 21: Total reward = 12.0\n",
            "Episode 22: Total reward = 11.0\n",
            "Episode 23: Total reward = 11.0\n",
            "Episode 24: Total reward = 14.0\n",
            "Episode 25: Total reward = 73.0\n",
            "Episode 26: Total reward = 21.0\n",
            "Episode 27: Total reward = 59.0\n",
            "Episode 28: Total reward = 84.0\n",
            "Episode 29: Total reward = 145.0\n",
            "Episode 30: Total reward = 47.0\n",
            "Episode 31: Total reward = 128.0\n",
            "Episode 32: Total reward = 63.0\n",
            "Episode 33: Total reward = 73.0\n",
            "Episode 34: Total reward = 43.0\n",
            "Episode 35: Total reward = 47.0\n",
            "Episode 36: Total reward = 200.0\n",
            "Episode 37: Total reward = 114.0\n",
            "Episode 38: Total reward = 66.0\n",
            "Episode 39: Total reward = 40.0\n",
            "Episode 40: Total reward = 54.0\n",
            "Episode 41: Total reward = 63.0\n",
            "Episode 42: Total reward = 48.0\n",
            "Episode 43: Total reward = 46.0\n",
            "Episode 44: Total reward = 39.0\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-4-57e36907c025>\u001b[0m in \u001b[0;36m<cell line: 10>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     17\u001b[0m         \u001b[0mtotal_reward\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mreward\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m         \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mremember\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maction\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreward\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnext_state\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 19\u001b[0;31m         \u001b[0magent\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreplay\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mdone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-1-43fcfa326021>\u001b[0m in \u001b[0;36mreplay\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     63\u001b[0m         \u001b[0mnext_q_values\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mq_network\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext_states\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkeepdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     64\u001b[0m         \u001b[0mtarget_q_values\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrewards\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0mdones\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgamma\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mnext_q_values\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 65\u001b[0;31m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq_values\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_q_values\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     66\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     67\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1492\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1493\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1494\u001b[0;31m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1495\u001b[0m         \u001b[0mforward_call\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_tracing_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1496\u001b[0m         \u001b[0;31m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ]
    }
  ]
}