gym_and_dqn.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "gym and dqn.ipynb",
      "private_outputs": true,
      "provenance": [],
      "authorship_tag": "ABX9TyNa3YsrfN8xAunRfmQxiQJI",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tigerneil/AGI-GAME/blob/master/gym_and_dqn.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o3jxY0liKlWA"
      },
      "source": [
        "%matplotlib inline\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import gym"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mMpauUr4Kvgm"
      },
      "source": [
        "all_envs = gym.envs.registry.all()\n",
        "env_ids = [env.id for env in all_envs]\n",
        "\n",
        "print(f'There are {len(env_ids)} gym environments. Such as {env_ids[:12]}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XWWnkGHSLg4g"
      },
      "source": [
        "# discrete action space environment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TvCivqg3LEfb"
      },
      "source": [
        "env = gym.make('CartPole-v1')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-pNqsl9lLWtl"
      },
      "source": [
        "print('observation space is:', env.observation_space)\n",
        "\n",
        "print('is observation space discrete?', isinstance(env.observation_space, gym.spaces.Discrete))\n",
        "print('is observation space continuous?', isinstance(env.observation_space, gym.spaces.Box))\n",
        "\n",
        "print('observation space shape:', env.observation_space.shape)\n",
        "\n",
        "print('observation space high values?', env.observation_space.high)\n",
        "print('observation space low values?', env.observation_space.low)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t32Hvk7aMQia"
      },
      "source": [
        "print('action space is:', env.action_space)\n",
        "\n",
        "print('is action space discrete?', isinstance(env.action_space, gym.spaces.Discrete))\n",
        "print('is action space continuous?', isinstance(env.action_space, gym.spaces.Box))\n",
        "\n",
        "print('action space shape:', env.action_space.n)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eY_1IvFTMgZI"
      },
      "source": [
        "print(env.spec.max_episode_steps)\n",
        "print(env.spec.reward_threshold)\n",
        "print(env.spec.nondeterministic)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zj1SvvwoMs8u"
      },
      "source": [
        "env = gym.make('Pendulum-v0')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "foi9Fd2MMzxf"
      },
      "source": [
        "print('observation space is:', env.observation_space)\n",
        "\n",
        "print('is observation space discrete?', isinstance(env.observation_space, gym.spaces.Discrete))\n",
        "print('is observation space continuous?', isinstance(env.observation_space, gym.spaces.Box))\n",
        "\n",
        "print('observation space shape:', env.observation_space.shape)\n",
        "\n",
        "print('observation space high values?', env.observation_space.high)\n",
        "print('observation space low values?', env.observation_space.low)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uvg0tKPeM2bv"
      },
      "source": [
        "print('action space is:', env.action_space)\n",
        "\n",
        "print('is action space discrete?', isinstance(env.action_space, gym.spaces.Discrete))\n",
        "print('is action space continuous?', isinstance(env.action_space, gym.spaces.Box))\n",
        "\n",
        "print('action space shape:', env.action_space.shape)\n",
        "\n",
        "print('action space high values?', env.action_space.high)\n",
        "print('action space low values?', env.action_space.low)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rCyCRvsXNBiF"
      },
      "source": [
        "\n",
        "print(env.spec.max_episode_steps)\n",
        "print(env.spec.reward_threshold)\n",
        "print(env.spec.nondeterministic)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZwR04h8KNEl_"
      },
      "source": [
        "env = gym.make('Freeway-v4')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6QyVvbOBNHth"
      },
      "source": [
        "pip install atari-py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UBF04z9hNObb"
      },
      "source": [
        "!wget http://www.atarimania.com/roms/Roms.rar"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_LbAUPn4Nb8J"
      },
      "source": [
        "!unrar Roms.rar"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q8ZaHOZmNjnd"
      },
      "source": [
        "ls"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u2fD-EWlNr9z"
      },
      "source": [
        "!unrar e -r Roms.rar"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VlkGJsAFNvRZ"
      },
      "source": [
        "ls"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QZ9THRrINxqK"
      },
      "source": [
        "ls 'HC ROMS.zip'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X8Ct7Q0lOC-9"
      },
      "source": [
        "!unzip ROMS.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eptlQar6OGgi"
      },
      "source": [
        "!python -m atari_py.import_roms ./ROMS\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nl7BzApsOKZi"
      },
      "source": [
        "print('observation space is:', env.observation_space)\n",
        "print('is observation space discrete?', isinstance(env.observation_space, gym.spaces.Discrete))\n",
        "print('is observation space continuous?', isinstance(env.observation_space, gym.spaces.Box))\n",
        "print('observation space shape:', env.observation_space.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sp16aNP3OVzz"
      },
      "source": [
        "print('action space is:', env.action_space)\n",
        "print('action space shape:', env.action_space.n)\n",
        "print('is action space discrete?', isinstance(env.action_space, gym.spaces.Discrete))\n",
        "print('is action space continuous?', isinstance(env.action_space, gym.spaces.Box))\n",
        "print('action meanings:', env.unwrapped.get_action_meanings())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z9bZ4F0OOa9Q"
      },
      "source": [
        "print(env.spec.max_episode_steps)\n",
        "print(env.spec.reward_threshold)\n",
        "print(env.spec.nondeterministic)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "roepBYqaOdxf"
      },
      "source": [
        "\n",
        "class ClipRewardEnv(gym.RewardWrapper):\n",
        "    def __init__(self, env):\n",
        "        gym.RewardWrapper.__init__(self, env)\n",
        "\n",
        "    def reward(self, reward):\n",
        "        \"\"\"Bin reward to {+1, 0, -1} by its sign.\"\"\"\n",
        "        return np.sign(reward)\n",
        "    \n",
        "class ScaledFloatFrame(gym.ObservationWrapper):\n",
        "    def __init__(self, env):\n",
        "        gym.ObservationWrapper.__init__(self, env)\n",
        "        self.observation_space = gym.spaces.Box(low=0, high=1, shape=env.observation_space.shape, dtype=np.float32)\n",
        "\n",
        "    def observation(self, observation):\n",
        "        return np.array(observation).astype(np.float32) / 255.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6ESImwN8Og9L"
      },
      "source": [
        "env = gym.make('Pong-v0')\n",
        "\n",
        "env = ClipRewardEnv(env)\n",
        "\n",
        "env = ScaledFloatFrame(env)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4_7RULLZOjB8"
      },
      "source": [
        "\n",
        "env = gym.make('CartPole-v1')\n",
        "\n",
        "env.seed(1234)\n",
        "\n",
        "state = env.reset()\n",
        "\n",
        "print('state type:', type(state))\n",
        "print('state shape:', state.shape)\n",
        "print('state:', state)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xkh84anLOl0k"
      },
      "source": [
        "action = env.action_space.sample() #select random action, uniformly between high and low for continuous\n",
        "\n",
        "print('selected action:', action)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nuhr2tYSOpNg"
      },
      "source": [
        "\n",
        "state, reward, done, info = env.step(action) #perform action on environment\n",
        "\n",
        "print('state:', state)\n",
        "print('reward:', reward)\n",
        "print('done:', done)\n",
        "print('info:', info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n5y2UbH8Ov2J"
      },
      "source": [
        "plt.imshow(state);\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WKTw32dGO4Sg"
      },
      "source": [
        "\n",
        "env = gym.make('FreewayNoFrameskip-v4')\n",
        "\n",
        "env.seed(1234)\n",
        "\n",
        "state = env.reset()\n",
        "\n",
        "print('state type:', type(state))\n",
        "print('state shape:', state.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x3ChrFGrPNN9"
      },
      "source": [
        "plt.imshow(state);\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UGxKpLARPPx8"
      },
      "source": [
        "action = env.action_space.sample() #select random action, uniformly between high and low for continuous\n",
        "\n",
        "print('selected action:', action)\n",
        "print('action meaning:', env.unwrapped.get_action_meanings()[action])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jaOfa7dVPTSj"
      },
      "source": [
        "state, reward, done, info = env.step(action) #perform action on environment\n",
        "\n",
        "print('reward:', reward)\n",
        "print('done:', done)\n",
        "print('info:', info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z2dgmF1CPWD2"
      },
      "source": [
        "plt.imshow(state);\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bOXRKSL8PX2R"
      },
      "source": [
        "\n",
        "up_action = env.unwrapped.get_action_meanings().index('UP')\n",
        "\n",
        "for i in range(50):\n",
        "    state, reward, done, info = env.step(up_action) #presses up 10 times\n",
        "\n",
        "plt.imshow(state);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9Ce9ToahPdpE"
      },
      "source": [
        "env = gym.make('SpaceInvadersNoFrameskip-v4')\n",
        "\n",
        "env.seed(1234)\n",
        "\n",
        "n_episodes = 10\n",
        "\n",
        "for episode in range(n_episodes):\n",
        "    \n",
        "    episode_reward = 0\n",
        "    done = False\n",
        "    state = env.reset()\n",
        "    \n",
        "    while not done:\n",
        "        \n",
        "        action = env.action_space.sample()\n",
        "        \n",
        "        state, reward, done, _ = env.step(action)\n",
        "\n",
        "        plt.imshow(state)\n",
        "        \n",
        "        episode_reward += reward\n",
        "        \n",
        "    print(f'episode: {episode+1}, reward: {episode_reward}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-u5Na3loPlpA"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torch.distributions as distributions\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import gym"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gNxgSWqFYkNe"
      },
      "source": [
        "train_env = gym.make('CartPole-v1')\n",
        "test_env = gym.make('CartPole-v1')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ae9NXtbtYnUM"
      },
      "source": [
        "SEED = 1234\n",
        "\n",
        "train_env.seed(SEED);\n",
        "test_env.seed(SEED+1);\n",
        "np.random.seed(SEED);\n",
        "torch.manual_seed(SEED);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d66RGJ0OYpwv"
      },
      "source": [
        "class MLP(nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.5):\n",
        "        super().__init__()\n",
        "\n",
        "        self.fc_1 = nn.Linear(input_dim, hidden_dim)\n",
        "        self.fc_2 = nn.Linear(hidden_dim, output_dim)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.fc_1(x)\n",
        "        x = self.dropout(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.fc_2(x)\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tmAycmO9ZUx6"
      },
      "source": [
        "INPUT_DIM = train_env.observation_space.shape[0]\n",
        "HIDDEN_DIM = 128\n",
        "OUTPUT_DIM = train_env.action_space.n\n",
        "\n",
        "policy = MLP(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sj-qJ2-iZnya"
      },
      "source": [
        "print(policy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SsUMPWo5Zs-d"
      },
      "source": [
        "def init_weights(m):\n",
        "    if type(m) == nn.Linear:\n",
        "        torch.nn.init.xavier_normal_(m.weight)\n",
        "        m.bias.data.fill_(0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LJVYYZMnZ9st"
      },
      "source": [
        "policy.apply(init_weights)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q18ctEmqaE_X"
      },
      "source": [
        "LEARNING_RATE = 0.01\n",
        "\n",
        "optimizer = optim.Adam(policy.parameters(), lr = LEARNING_RATE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kjAVdlANaaH4"
      },
      "source": [
        "def train(env, policy, optimizer, discount_factor):\n",
        "    policy.train()\n",
        "\n",
        "    log_prob_actions = []\n",
        "    rewards = []\n",
        "    done = False\n",
        "    episode_reward = 0\n",
        "\n",
        "    state = env.reset()\n",
        "\n",
        "    while not done:\n",
        "        state = torch.FloatTensor(state).unsqueeze(0).cuda()\n",
        "\n",
        "        action_pred = policy(state).cuda()\n",
        "\n",
        "        action_prob = F.softmax(action_pred, dim = -1).cuda()\n",
        "\n",
        "        dist = distributions.Categorical(action_prob)\n",
        "\n",
        "        action = dist.sample().cuda()\n",
        "\n",
        "        log_prob_action = dist.log_prob(action).cuda()\n",
        "\n",
        "        state, reward, done, _ = env.step(action.item())\n",
        "\n",
        "        log_prob_actions.append(log_prob_action)\n",
        "\n",
        "        rewards.append(reward)\n",
        "\n",
        "        episode_reward += reward\n",
        "\n",
        "    log_prob_actions = torch.cat(log_prob_actions)\n",
        "\n",
        "    returns = calculate_returns(rewards, discount_factor)\n",
        "\n",
        "    loss = update_policy(returns, log_prob_actions, optimizer)\n",
        "\n",
        "    return loss, episode_reward"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g2NcXkD-bts0"
      },
      "source": [
        "def calculate_returns(rewards, discount_factor, normalize=True):\n",
        "    returns = []\n",
        "    R = 0\n",
        "\n",
        "    for r in reversed(rewards):\n",
        "        R = r + R * discount_factor\n",
        "        returns.insert(0, R)\n",
        "    \n",
        "    returns = torch.tensor(returns)\n",
        "\n",
        "    if normalize:\n",
        "        returns = (returns - returns.mean()) / returns.std()\n",
        "\n",
        "    return returns"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wbuNv7FXcNSm"
      },
      "source": [
        "def update_policy(returns, log_prob_actions, optimizer):\n",
        "    returns = returns.detach().cuda()\n",
        "    loss = - (returns * log_prob_actions).sum()\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "    return loss.item()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YaHeOjBHchts"
      },
      "source": [
        "def evaluate(env, policy):\n",
        "    policy.eval()\n",
        "\n",
        "    done = False\n",
        "    episode_reward = 0\n",
        "\n",
        "    state = env.reset()\n",
        "\n",
        "    while not done:\n",
        "        state = torch.FloatTensor(state).unsqueeze(0).cuda()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            action_pred = policy(state).cuda()\n",
        "\n",
        "            action_prob = F.softmax(action_pred, dim = -1).cuda()\n",
        "\n",
        "        action = torch.argmax(action_prob, dim = -1).cuda()\n",
        "\n",
        "        state, reward, done, _ = env.step(action.item())\n",
        "\n",
        "        episode_reward += reward\n",
        "\n",
        "    return episode_reward"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2vRRyYc3dK1c"
      },
      "source": [
        "MAX_EPISODES = 500\n",
        "DISCOUNT_FACTOR = 0.99\n",
        "N_TRIALS = 25\n",
        "REWARD_THRESHOLD =1000\n",
        "PRINT_EVERY = 10\n",
        "\n",
        "train_rewards = []\n",
        "test_rewards = []\n",
        "\n",
        "for episode in range(1, MAX_EPISODES+1):\n",
        "    loss, train_reward = train(train_env, policy, optimizer, DISCOUNT_FACTOR)\n",
        "\n",
        "    test_reward = evaluate(test_env, policy)\n",
        "\n",
        "    train_rewards.append(train_reward)\n",
        "    test_rewards.append(test_reward)\n",
        "\n",
        "    mean_train_rewards = np.mean(train_rewards[-N_TRIALS:])\n",
        "    mean_test_rewards = np.mean(test_rewards[-N_TRIALS:])\n",
        "\n",
        "    if episode % PRINT_EVERY == 0:\n",
        "        print(f'| Episode: {episode:3} | Mean Train Rewards: {mean_train_rewards:5.1f} | Mean Test Rewards: {mean_test_rewards:5.1f} |')\n",
        "\n",
        "    if mean_test_rewards >= REWARD_THRESHOLD:\n",
        "        print(f'Reached reward threshold in {episode} episodes')\n",
        "\n",
        "        break"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2-uz5jVWerAB"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}