reinforce_AC_CartPole_torch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "reinforce-AC-CartPole-torch.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "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.7.4"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/oroojlooy/RL_pytorch/blob/master/reinforce_AC_CartPole_torch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eJqIx--gMVhT"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qjTnUfcUA_Bh"
      },
      "source": [
        "# !pip install --upgrade gym\n",
        "# !pip install --upgrade torch"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cmqNcyIXAT8Y"
      },
      "source": [
        "import argparse\n",
        "import gym\n",
        "import numpy as np\n",
        "from itertools import count\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.distributions import Categorical\n",
        "\n",
        "import time"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k7Y73Kd4AT8d"
      },
      "source": [
        "# parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rwbldOLQAT8f"
      },
      "source": [
        "parser = argparse.ArgumentParser(description='PyTorch REINFORCE example')\n",
        "parser.add_argument('--gamma', type=float, default=0.99, metavar='G',\n",
        "                    help='discount factor (default: 0.99)')\n",
        "parser.add_argument('--seed', type=int, default=543, metavar='N',\n",
        "                    help='random seed (default: 543)')\n",
        "parser.add_argument('--render', action='store_true',\n",
        "                    help='render the environment')\n",
        "parser.add_argument('--log-interval', type=int, default=200, metavar='N',\n",
        "                    help='interval between training status logs (default: 10)')\n",
        "args, unknown = parser.parse_known_args()\n"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S7tCuXp-AT8i"
      },
      "source": [
        "# create an environment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fcX1dN16AT8j",
        "outputId": "bca8df6e-372a-48ea-e9c5-180da8ea9067"
      },
      "source": [
        "# this environment has env.reset() and end.step() functions\n",
        "env = gym.make('CartPole-v0')\n",
        "env.seed(args.seed)\n",
        "torch.manual_seed(args.seed)\n"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7f5b75b60060>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-y254LZFAT8p",
        "outputId": "3a4a8af7-e6a3-46c4-bb4d-59c98735a2e0"
      },
      "source": [
        "env.reset()\n",
        "# env.render()\n",
        "env.step(1)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(array([ 0.00313956,  0.2222067 , -0.01409669, -0.31700878]), 1.0, False, {})"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X7jaeYIlAT8u"
      },
      "source": [
        "# create actor network"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y8K7leovBfIa"
      },
      "source": [
        ""
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MnGQtBAEAT8v"
      },
      "source": [
        "class Actor(nn.Module):\n",
        "    # this class defines a policy network with two layer NN\n",
        "    def __init__(self):\n",
        "        super(Actor, self).__init__()\n",
        "        self.affine1 = nn.Linear(4, 30)\n",
        "        self.affine2 = nn.Linear(30, 2)\n",
        "\n",
        "    def forward(self, x):\n",
        "        ''' do the forward pass and return a probability over actions\n",
        "        Input:\n",
        "                x: state\n",
        "        returns:\n",
        "                prob: a probability distribution\n",
        "        '''\n",
        "        \n",
        "        x = F.relu(self.affine1(x))\n",
        "        action_scores = self.affine2(x)\n",
        "        prob = F.softmax(action_scores, dim=1)\n",
        "        return prob"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QXQuyDXhAT8z"
      },
      "source": [
        "# Critic network"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5TlyZZTQAT8z"
      },
      "source": [
        "class Critic(nn.Module):\n",
        "    # this class defines a policy network with two layer NN\n",
        "    def __init__(self, in_d=4):\n",
        "        super(Critic, self).__init__()\n",
        "        self.affine1 = nn.Linear(in_d, 30)\n",
        "        self.affine2 = nn.Linear(30, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        ''' do the forward pass and return a probability over actions\n",
        "        Input:\n",
        "                x: state\n",
        "        returns:\n",
        "                v: value of being at x\n",
        "        '''\n",
        "        \n",
        "        x = F.relu(self.affine1(x))\n",
        "        v = self.affine2(x).squeeze()\n",
        "        return v\n",
        "\n",
        "class QCritic(nn.Module):\n",
        "    # this class defines a policy network with two layer NN\n",
        "    def __init__(self, in_d=4):\n",
        "        super(Critic, self).__init__()\n",
        "        self.affine1 = nn.Linear(in_d, 30)\n",
        "        self.affine2 = nn.Linear(30, 2)\n",
        "\n",
        "    def forward(self, x):\n",
        "        ''' do the forward pass and return a probability over actions\n",
        "        Input:\n",
        "                x: state\n",
        "        returns:\n",
        "                v: value of being at x\n",
        "        '''\n",
        "        \n",
        "        x = F.relu(self.affine1(x))\n",
        "        v = self.affine2(x).squeeze()\n",
        "        return v\n"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Vs-9MAYAT83"
      },
      "source": [
        "# Create networks and optimizers\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-__a1C3aAT84"
      },
      "source": [
        ""
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "az8jMJnnDoau"
      },
      "source": [
        ""
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ATtRw2qSAT8-"
      },
      "source": [
        "# rollout funtion"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lBWDaywwAT9A"
      },
      "source": [
        "def select_action(state):\n",
        "    # this function selects stochastic actions based on the policy probabilities\n",
        "    state = torch.from_numpy(state).float().unsqueeze(0)\n",
        "    probs = actor(state)\n",
        "    m = Categorical(probs)\n",
        "    action = m.sample()\n",
        "    log_prob = m.log_prob(action)\n",
        "    \n",
        "    return action.item(), log_prob"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tfowBl_lBsUW"
      },
      "source": [
        "# s=env.reset()\n",
        "# # select_action(s)\n",
        "# state = torch.from_numpy(s).float().unsqueeze(0)\n",
        "# probs = actor(state)\n",
        "# m = Categorical(probs)\n",
        "# action = m.sample()\n",
        "# log_prob = m.log_prob(action)\n",
        "# print(probs, m, action, log_prob)\n",
        "# # (tensor([[0.5007, 0.4993]], grad_fn=<SoftmaxBackward>), Categorical(), tensor([1]), tensor([-0.6946], grad_fn=<SqueezeBackward1>))\n",
        "# np.log(0.4993)\n",
        "# # -0.6945481614755734"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w_RYkLaKCPxR"
      },
      "source": [
        ""
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rZ05O7LkuTjs"
      },
      "source": [
        ""
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TAzOnWWKAT9F"
      },
      "source": [
        "def rollout(render=False,pause=.2):\n",
        "    states = []\n",
        "    rewards = []\n",
        "    log_probs = []\n",
        "    next_states = []\n",
        "    mask = []\n",
        "    # play an episode\n",
        "    state = env.reset()\n",
        "    while True:  # Don't infinite loop while learning\n",
        "        # select an action\n",
        "        action, log_prob = select_action(state)\n",
        "        states.append(list(state))\n",
        "        log_probs.append(log_prob[0])\n",
        "        \n",
        "        # take the action and move to next state\n",
        "        next_state, reward, done, _ = env.step(action)\n",
        "        rewards.append(reward)\n",
        "        next_states.append(next_state)\n",
        "        if render:\n",
        "            env.render()\n",
        "            time.sleep(pause)\n",
        "        if done:\n",
        "            mask.append(1)\n",
        "            break\n",
        "        mask.append(0)\n",
        "        state = next_state\n",
        "            \n",
        "    return states, rewards, log_probs, next_states, mask"
      ],
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NSZPgcmcAT9I"
      },
      "source": [
        "# train function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Bq_H6jcUAT9J"
      },
      "source": [
        "def reinforce_train_step(states,rewards,log_probs, critic):\n",
        "    \n",
        "    R = 0\n",
        "    P = 0\n",
        "    rewards_path = []\n",
        "    log_probs_paths = []\n",
        "    for i in reversed(range(len(rewards))):\n",
        "        R = rewards[i] + args.gamma * R\n",
        "        rewards_path.insert(0, R) \n",
        "        \n",
        "        P = log_probs[i] + P\n",
        "        log_probs_paths.insert(0, P) \n",
        "\n",
        "    rewards_path = torch.tensor(rewards_path)\n",
        "    rewards_path = (rewards_path - rewards_path.mean()) / (rewards_path.std() + 1e-8)\n",
        "    log_probs_paths = torch.stack(log_probs_paths)\n",
        "    \n",
        "#     print(rewards_path,log_probs_paths)\n",
        "    value = critic(torch.tensor(states, dtype=torch.float32))\n",
        "\n",
        "    # take a backward step for actor\n",
        "    # This is based on Pytorch implementation of REINFORCE \n",
        "    actor_loss = -torch.mean(((rewards_path - value.detach()) * torch.stack(log_probs)))\n",
        "    # This is based on the formual which Levine obtains for REINFORCE algorithm\n",
        "    # actor_loss = -torch.mean(((rewards_path - value.detach()) * log_probs_paths))\n",
        "    actor_optim.zero_grad()\n",
        "    actor_loss.backward()\n",
        "    actor_optim.step()\n",
        "\n",
        "    # take a backward step for critic\n",
        "    loss_fn = torch.nn.MSELoss()\n",
        "    critic_loss = loss_fn(value, rewards_path)\n",
        "    critic_optim.zero_grad()\n",
        "    critic_loss.backward()\n",
        "    critic_optim.step()\n"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hfAh0rSMaDl4"
      },
      "source": [
        "# Actor Critic \n",
        "def ac_train_step(states, rewards, log_probs, next_states, done, critic):\n",
        "# states, rewards, log_probs, next_states = rollout()\n",
        "# if 1 == 1:\n",
        "\n",
        "    P = 0\n",
        "    rewards_path = []\n",
        "    log_probs_paths = []\n",
        "    qvalue = critic(torch.tensor(states, dtype=torch.float32))\n",
        "    next_qvalue = critic(torch.tensor(next_states, dtype=torch.float32))\n",
        "    next_qvalue = next_qvalue.detach().cpu().numpy()\n",
        "    target = []\n",
        "    for i in range(len(rewards)):\n",
        "        target.append(rewards[i] + args.gamma * next_qvalue[i]*(1-done[i]))\n",
        "        \n",
        "    target = torch.tensor(target)\n",
        "\n",
        "    qvalue = (qvalue - torch.mean(qvalue))/(torch.std(qvalue) + 1e-8)\n",
        "    # take a backward step for actor\n",
        "    # This is based on the commonly used AC algorithm  \n",
        "    actor_loss = -torch.mean(qvalue.detach() * torch.stack(log_probs))\n",
        "    actor_optim.zero_grad()\n",
        "    actor_loss.backward()\n",
        "    actor_optim.step()\n",
        "\n",
        "    # take a backward step for critic\n",
        "    loss_fn = torch.nn.MSELoss()\n",
        "    critic_loss = loss_fn(qvalue, torch.tensor(target, dtype=torch.float32))\n",
        "    critic_optim.zero_grad()\n",
        "    critic_loss.backward()\n",
        "    critic_optim.step()\n"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u69K431OFQf-"
      },
      "source": [
        "def rtg_reinforce_train_step(states,rewards_,log_probs, critic):\r\n",
        "    \r\n",
        "    R = 0\r\n",
        "    G = 0\r\n",
        "    log_probs_grads = []\r\n",
        "    rewards = []\r\n",
        "#     print(rewards_path,log_probs_paths)\r\n",
        "    value = critic(torch.tensor(states, dtype=torch.float32))\r\n",
        "    for i in reversed(range(len(rewards_))):\r\n",
        "        R = rewards_[i] + args.gamma * R\r\n",
        "        rewards.insert(0, R)\r\n",
        "        \r\n",
        "    rewards = torch.tensor(rewards)\r\n",
        "    rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-8)\r\n",
        "    value_detached = value.detach()\r\n",
        "    for i in reversed(range(len(rewards_))):\r\n",
        "        G += log_probs[i]*(rewards[i] - value_detached[i]) \r\n",
        "        log_probs_grads.insert(0, G) \r\n",
        "\r\n",
        "    log_probs_grads = torch.stack(log_probs_grads)\r\n",
        "    \r\n",
        "    # take a backward step for actor\r\n",
        "    # This is based on Pytorch implementation of REINFORCE \r\n",
        "    actor_loss = -torch.mean(log_probs_grads)\r\n",
        "    actor_optim.zero_grad()\r\n",
        "    actor_loss.backward()\r\n",
        "    actor_optim.step()\r\n",
        "\r\n",
        "    # take a backward step for critic\r\n",
        "    loss_fn = torch.nn.MSELoss()\r\n",
        "    critic_loss = loss_fn(value, rewards)\r\n",
        "    critic_optim.zero_grad()\r\n",
        "    critic_loss.backward()\r\n",
        "    critic_optim.step()\r\n"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c-asw4wcRh-G"
      },
      "source": [
        "# Advantage Actor Critic  (single worker A2C)\n",
        "def aac_train_step(states, rewards, log_probs, next_states, done, critic):\n",
        "# states, rewards, log_probs, next_states = rollout()\n",
        "# if 1 == 1:\n",
        "    P = 0\n",
        "    rewards_path = []\n",
        "    log_probs_paths = []\n",
        "    value = critic(torch.tensor(states, dtype=torch.float32))\n",
        "    next_value = critic(torch.tensor(next_states, dtype=torch.float32))\n",
        "    next_value = next_value.detach().cpu().numpy()\n",
        "    advantage = []\n",
        "    target = []\n",
        "    for i in range(len(rewards)):\n",
        "        target.append(rewards[i] + args.gamma * next_value[i]*(1-done[i]))\n",
        "        advantage.append(target[i] - value[i].detach().cpu().numpy())\n",
        "        \n",
        "        P = log_probs[i] + P\n",
        "        log_probs_paths.insert(0, P) \n",
        "\n",
        "    advantage = torch.tensor(advantage)\n",
        "    target = torch.tensor(target)\n",
        "    cumulative_log_probs = torch.stack(log_probs_paths)\n",
        "    \n",
        "#     print(log_probs_paths)\n",
        "\n",
        "    # take a backward step for actor\n",
        "    # This is based on Pytorch implementation of AC \n",
        "    actor_loss = -torch.mean(advantage.detach() * torch.stack(log_probs))\n",
        "    # This is based on the formual which Levine obtains for AC algorithm\n",
        "    # actor_loss = -torch.mean(advantage.detach() * cumulative_log_probs)\n",
        "    actor_optim.zero_grad()\n",
        "    actor_loss.backward()\n",
        "    actor_optim.step()\n",
        "\n",
        "    # take a backward step for critic\n",
        "    loss_fn = torch.nn.MSELoss()\n",
        "    critic_loss = loss_fn(value, torch.tensor(target, dtype=torch.float32))\n",
        "    critic_optim.zero_grad()\n",
        "    critic_loss.backward()\n",
        "    critic_optim.step()\n"
      ],
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n4zw7GSYAT9M"
      },
      "source": [
        "# run training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2FGjCmdndDcj",
        "outputId": "68062ac0-a46e-4c25-bd81-29c89720987f"
      },
      "source": [
        "# create actor and critic network\n",
        "alg = 'acrtg' # 'reinforce'\n",
        "if alg == 'ac_q':\n",
        "  critic = QCritic()\n",
        "else:\n",
        "  critic = Critic()\n",
        "\n",
        "actor = Actor()\n",
        "\n",
        "# create optimizers\n",
        "actor_optim = optim.Adam(actor.parameters(), lr=1e-3)\n",
        "critic_optim = optim.Adam(critic.parameters(), lr=1e-3)\n",
        "\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)\n",
        "print()\n",
        "\n",
        "#Additional Info when using cuda\n",
        "if device.type == 'cuda':\n",
        "    print(torch.cuda.get_device_name(0))\n",
        "    print('Memory Usage:')\n",
        "    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')\n",
        "    print('Cached:   ', round(torch.cuda.memory_reserved(0)/1024**3,1), 'GB')"
      ],
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using device: cpu\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Pxqb4tH4AT9N",
        "scrolled": true,
        "outputId": "46ed585d-668d-4f55-e2bd-a836bdd2cbbd"
      },
      "source": [
        "running_reward = 10.\n",
        "for i_episode in range(2000):\n",
        "    states, rewards, log_probs, next_states, done = rollout()\n",
        "    t = len(rewards)\n",
        "    running_reward = 0.9*running_reward + 0.1*t\n",
        "    if alg == 'reinforce':\n",
        "        reinforce_train_step(states, rewards, log_probs, critic)\n",
        "    elif alg == 'ac':\n",
        "        ac_train_step(states, rewards, log_probs, next_states, done, critic)\n",
        "    elif alg == 'acrtg':\n",
        "        rtg_reinforce_train_step(states, rewards, log_probs, critic)\n",
        "    elif alg == 'aac':\n",
        "        aac_train_step(states, rewards, log_probs, next_states, done, critic)\n",
        "    if i_episode % args.log_interval == 0:\n",
        "        print('Episode {}\\tLast length: {:5d}\\tAverage length: {:.2f}'.format(\n",
        "            i_episode, t, running_reward))\n",
        "    if running_reward > env.spec.reward_threshold:\n",
        "        print(\"Solved! Running reward is now {} and \"\n",
        "              \"the last episode runs to {} time steps!\".format(running_reward, t))\n",
        "        break"
      ],
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Episode 0\tLast length:    10\tAverage length: 10.00\n",
            "Episode 200\tLast length:     9\tAverage length: 19.39\n",
            "Episode 400\tLast length:    26\tAverage length: 44.57\n",
            "Episode 600\tLast length:   159\tAverage length: 170.63\n",
            "Solved! Running reward is now 195.20186849889512 and the last episode runs to 200 time steps!\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YDo7p8QBczwa"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6jogUgc5cYyf",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 231
        },
        "outputId": "ec33c59e-dd4a-4304-9fe3-55770969acd7"
      },
      "source": [
        "running_reward = 10\n",
        "for i_episode in range(1000):\n",
        "    states, rewards, log_probs = rollout()\n",
        "    t = len(rewards)\n",
        "    running_reward = running_reward * 0.9 +  t * 0.1\n",
        "    ac_train_step(states, rewards, log_probs)\n",
        "    if i_episode % args.log_interval == 0:\n",
        "        print('Episode {}\\tLast length: {:5d}\\tAverage length: {:.2f}'.format(\n",
        "            i_episode, t, running_reward))\n",
        "    if running_reward > env.spec.reward_threshold:\n",
        "        print(\"Solved! Running reward is now {} and \"\n",
        "              \"the last episode runs to {} time steps!\".format(running_reward, t))\n",
        "        break"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-27-0c8618d0dda5>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mrunning_reward\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi_episode\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1000\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----> 3\u001b[0;31m     \u001b[0mstates\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrewards\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlog_probs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrollout\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      4\u001b[0m     \u001b[0mt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrewards\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mrunning_reward\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrunning_reward\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;36m0.9\u001b[0m \u001b[0;34m+\u001b[0m  \u001b[0mt\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0;36m0.1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mValueError\u001b[0m: too many values to unpack (expected 3)"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6T_v8LrRcYuU"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cPxtYu3DqzSl"
      },
      "source": [
        "# for var_name in actor_optim.state_dict():\n",
        "#     print(var_name, \"\\t\", actor_optim.state_dict()[var_name])\n",
        "\n",
        "d=actor_optim.state_dict()\n",
        "# d\n",
        "optimizer_state= actor_optim.state # holds all the information about the gradients, square of gradients and past gradients used in Adam \n",
        "\n",
        "te=d['state'][140575074157424]['exp_avg']\n",
        "te.shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "erXIKeZZrGq2"
      },
      "source": [
        "c=torch.nn.Conv1d(10,12,3)\n",
        "c=torch.nn.Conv2d(10,12,3)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V95uZz_mrGt3"
      },
      "source": [
        "for i in c.parameters():\n",
        "#   print (i['weight'])\n",
        "#   print (dir(i))\n",
        "  print (i.shape)\n",
        "  print (i.numel())\n",
        "#   print (i.count())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5k_5eaXAT9T"
      },
      "source": [
        "# render optimal policy\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uchezPDiAT9V"
      },
      "source": [
        "a = rollout(True,pause=.05)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "42D2h_DqAT9Y"
      },
      "source": [
        "env.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1J149hSkAT9b"
      },
      "source": [
        "z=[5,7,9,12]\n",
        "for b in [1,2,10,30]:\n",
        "  print (\"------\",b)\n",
        "  s=sum([np.exp(i*b) for i in z])\n",
        "  for i in z:\n",
        "    print (np.exp(i*b)/s)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vw8EBC7LaEJ-"
      },
      "source": [
        "import torch\n",
        "\n",
        "batch_size = 5\n",
        "nb_digits = 10\n",
        "# Dummy input that HAS to be 2D for the scatter (you can use view(-1,1) if needed)\n",
        "y = torch.LongTensor(batch_size,1).random_() % nb_digits\n",
        "# One hot encoding buffer that you create out of the loop and just keep reusing\n",
        "y_onehot = torch.FloatTensor(batch_size, nb_digits)\n",
        "# In your for loop\n",
        "y_onehot.zero_()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BAnWaDiWpsb9"
      },
      "source": [
        "y_onehot.scatter_(1, y, 1)\n",
        "\n",
        "print(y)\n",
        "print(y_onehot)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mwKVOFVZptyq"
      },
      "source": [
        "d=torch.chunk(y_onehot,10,1)\n",
        "d[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aCekDr840unT"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}