minDQN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Untitled5.ipynb",
      "provenance": [],
      "private_outputs": true,
      "authorship_tag": "ABX9TyMwTRQM9UDgG+xw9K/B+sYW",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mswang12/minDQN/blob/main/minDQN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ziFAB5fMbwCB"
      },
      "source": [
        "# A Minimal Deep Q-Network\n",
        "We'll be showing how to code a minimal Deep Q-Network to solve the CartPole environment.\n",
        "\n",
        "## Step 1. Import libraries and setup the environment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_yfWkhBMbjR8"
      },
      "source": [
        "import gym\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "from tensorflow import keras\n",
        "\n",
        "from collections import deque\n",
        "import time\n",
        "import random\n",
        "\n",
        "RANDOM_SEED = 5\n",
        "tf.random.set_seed(RANDOM_SEED)\n",
        "\n",
        "env = gym.make('CartPole-v1')\n",
        "env.seed(RANDOM_SEED)\n",
        "np.random.seed(RANDOM_SEED)\n",
        "\n",
        "print(\"Action Space: {}\".format(env.action_space))\n",
        "print(\"State space: {}\".format(env.observation_space))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fNP9CoTTcDs3"
      },
      "source": [
        "## Step 2. Define the network architecture"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5NEKInS5bpjW"
      },
      "source": [
        "# An episode a full game\n",
        "train_episodes = 300\n",
        "test_episodes = 100\n",
        "\n",
        "def agent(state_shape, action_shape):\n",
        "    \"\"\" The agent maps X-states to Y-actions\n",
        "    e.g. The neural network output is [.1, .7, .05, 0.05, .05, .05]\n",
        "    The highest value 0.7 is the Q-Value.\n",
        "    The index of the highest action (0.7) is action #1.\n",
        "    \"\"\"\n",
        "    learning_rate = 0.001\n",
        "    init = tf.keras.initializers.HeUniform()\n",
        "    model = keras.Sequential()\n",
        "    model.add(keras.layers.Dense(24, input_shape=state_shape, activation='relu', kernel_initializer=init))\n",
        "    model.add(keras.layers.Dense(12, activation='relu', kernel_initializer=init))\n",
        "    model.add(keras.layers.Dense(action_shape, activation='linear', kernel_initializer=init))\n",
        "    model.compile(loss=tf.keras.losses.Huber(), optimizer=tf.keras.optimizers.Adam(lr=learning_rate), metrics=['accuracy'])\n",
        "    return model\n",
        "\n",
        "def get_qs(model, state, step):\n",
        "    return model.predict(state.reshape([1, state.shape[0]]))[0]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N47IymMecGmo"
      },
      "source": [
        "## Step 3. Define the train function using Experience Replay"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n2ZBvh3ZbscS"
      },
      "source": [
        "def train(env, replay_memory, model, target_model, done):\n",
        "    learning_rate = 0.7 # Learning rate\n",
        "    discount_factor = 0.618\n",
        "\n",
        "    MIN_REPLAY_SIZE = 1000\n",
        "    if len(replay_memory) < MIN_REPLAY_SIZE:\n",
        "        return\n",
        "\n",
        "    batch_size = 64 * 2\n",
        "    mini_batch = random.sample(replay_memory, batch_size)\n",
        "    current_states = np.array([encode_observation(transition[0], env.observation_space.shape) for transition in mini_batch])\n",
        "    current_qs_list = model.predict(current_states)\n",
        "    new_current_states = np.array([encode_observation(transition[3], env.observation_space.shape) for transition in mini_batch])\n",
        "    future_qs_list = target_model.predict(new_current_states)\n",
        "\n",
        "    X = []\n",
        "    Y = []\n",
        "    for index, (observation, action, reward, new_observation, done) in enumerate(mini_batch):\n",
        "        if not done:\n",
        "            max_future_q = reward + discount_factor * np.max(future_qs_list[index])\n",
        "        else:\n",
        "            max_future_q = reward\n",
        "\n",
        "        current_qs = current_qs_list[index]\n",
        "        current_qs[action] = (1 - learning_rate) * current_qs[action] + learning_rate * max_future_q\n",
        "\n",
        "        X.append(encode_observation(observation, env.observation_space.shape))\n",
        "        Y.append(current_qs)\n",
        "    model.fit(np.array(X), np.array(Y), batch_size=batch_size, verbose=0, shuffle=True)\n",
        "\n",
        "def encode_observation(observation, n_dims):\n",
        "    return observation"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9bxAPI27cQ-F"
      },
      "source": [
        "## Step 4. Run the Deep Q-Network Algorithm"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DOaFgqEVbuv-"
      },
      "source": [
        "def main():\n",
        "    epsilon = 1 # Epsilon-greedy algorithm in initialized at 1 meaning every step is random at the start\n",
        "    max_epsilon = 1 # You can't explore more than 100% of the time\n",
        "    min_epsilon = 0.01 # At a minimum, we'll always explore 1% of the time\n",
        "    decay = 0.01\n",
        "\n",
        "    # 1. Initialize the Target and Main models\n",
        "    # Main Model (updated every 4 steps)\n",
        "    model = agent(env.observation_space.shape, env.action_space.n)\n",
        "    # Target Model (updated every 100 steps)\n",
        "    target_model = agent(env.observation_space.shape, env.action_space.n)\n",
        "    target_model.set_weights(model.get_weights())\n",
        "\n",
        "    replay_memory = deque(maxlen=50_000)\n",
        "\n",
        "    target_update_counter = 0\n",
        "\n",
        "    # X = states, y = actions\n",
        "    X = []\n",
        "    y = []\n",
        "\n",
        "    steps_to_update_target_model = 0\n",
        "\n",
        "    for episode in range(train_episodes):\n",
        "        total_training_rewards = 0\n",
        "        observation = env.reset()\n",
        "        done = False\n",
        "        while not done:\n",
        "            steps_to_update_target_model += 1\n",
        "            #if True:\n",
        "            #    env.render()\n",
        "\n",
        "            random_number = np.random.rand()\n",
        "            # 2. Explore using the Epsilon Greedy Exploration Strategy\n",
        "            if random_number <= epsilon:\n",
        "                # Explore\n",
        "                action = env.action_space.sample()\n",
        "            else:\n",
        "                # Exploit best known action\n",
        "                # model dims are (batch, env.observation_space.n)\n",
        "                encoded = encode_observation(observation, env.observation_space.shape[0])\n",
        "                encoded_reshaped = encoded.reshape([1, encoded.shape[0]])\n",
        "                predicted = model.predict(encoded_reshaped).flatten()\n",
        "                action = np.argmax(predicted)\n",
        "            new_observation, reward, done, info = env.step(action)\n",
        "            replay_memory.append([observation, action, reward, new_observation, done])\n",
        "\n",
        "            # 3. Update the Main Network using the Bellman Equation\n",
        "            if steps_to_update_target_model % 4 == 0 or done:\n",
        "                train(env, replay_memory, model, target_model, done)\n",
        "\n",
        "            observation = new_observation\n",
        "            total_training_rewards += reward\n",
        "\n",
        "            if done:\n",
        "                print('Total training rewards: {} after n steps = {} with final reward = {}'.format(total_training_rewards, episode, reward))\n",
        "                total_training_rewards += 1\n",
        "\n",
        "                if steps_to_update_target_model >= 100:\n",
        "                    print('Copying main network weights to the target network weights')\n",
        "                    target_model.set_weights(model.get_weights())\n",
        "                    steps_to_update_target_model = 0\n",
        "                break\n",
        "\n",
        "        epsilon = min_epsilon + (max_epsilon - min_epsilon) * np.exp(-decay * episode)\n",
        "    env.close()\n",
        "\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    main()\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}