SpyTradingAgent.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SpyTradingAgent",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "private_outputs": true,
      "mount_file_id": "1EXh-2mp-AwE3OsMXYMgFGNCtenUG_T67",
      "authorship_tag": "ABX9TyPtBJZgQQqNEx/QbykvHIQR",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aCStandke/ReinforcementLearning/blob/main/SpyTradingAgent.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import warnings\n",
        "warnings.filterwarnings('ignore')"
      ],
      "metadata": {
        "id": "DEKeeEoUy2q4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Trading Using Reinforcment Learning\n",
        "\n",
        "\n",
        "This part implementation comes from chapter 8 from the  book \n",
        "[Deep Reinforcement Learning Hands-On - Second Edition by Maxim Lapan](https://www.amazon.com/Deep-Reinforcement-Learning-Hands-optimization/dp/1838826998/ref=asc_df_1838826998/?tag=hyprod-20&linkCode=df0&hvadid=416741343328&hvpos=&hvnetw=g&hvrand=7234438034400691228&hvpone=&hvptwo=&hvqmt=&hvdev=c&hvdvcmdl=&hvlocint=&hvlocphy=9008183&hvtargid=pla-871456510229&psc=1&tag=&ref=&adgrpid=93867144477&hvpone=&hvptwo=&hvadid=416741343328&hvpos=&hvnetw=g&hvrand=7234438034400691228&hvqmt=&hvdev=c&hvdvcmdl=&hvlocint=&hvlocphy=9008183&hvtargid=pla-871456510229)\n",
        "\n",
        " \n",
        "Namely, a RL agent has some observation of the market, and has to take an action to either buy, sell, or hold. If the agent buys before the price goes up, profit will be positive; otherwise, the agent will get a negative reward. The agent is tyring to obtain as much profit as possible in the trading environment. \n",
        "\n"
      ],
      "metadata": {
        "id": "iTCX_DxmyQLZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install ptan"
      ],
      "metadata": {
        "id": "i42g9KIlrVJL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tensorboardX"
      ],
      "metadata": {
        "id": "ntC5suZYsSvZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "import torch \n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.nn.utils as nn_utils\n",
        "import gym\n",
        "import gym.spaces\n",
        "from gym.utils import seeding\n",
        "from gym.envs.registration import EnvSpec\n",
        "import enum\n",
        "import glob\n",
        "import os\n",
        "import collections\n",
        "import csv\n",
        "import sys\n",
        "import time\n",
        "import ptan\n",
        "import torch.optim as optim\n",
        "from tensorboardX import SummaryWriter\n",
        "import matplotlib as mpl\n",
        "mpl.use(\"Agg\")\n",
        "import matplotlib.pyplot as plt"
      ],
      "metadata": {
        "id": "I7jj6N5k264p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Validation Environment"
      ],
      "metadata": {
        "id": "4IpU49YBeYqF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# env for validation (i.e. future market data for the agent to trade on)\n",
        "def validation_run(env, net, episodes=100, device=\"cpu\", epsilon=0.02, comission=0.1):\n",
        "    stats = {\n",
        "        'episode_reward': [],\n",
        "        'episode_steps': [],\n",
        "        'order_profits': [],\n",
        "        'order_steps': [],\n",
        "    }\n",
        "\n",
        "    for episode in range(episodes):\n",
        "        obs = env.reset()\n",
        "\n",
        "        total_reward = 0.0\n",
        "        position = None\n",
        "        position_steps = None\n",
        "        episode_steps = 0\n",
        "\n",
        "        while True:\n",
        "            obs_v = torch.tensor([obs]).to(device)\n",
        "            out_v = net(obs_v)\n",
        "\n",
        "            action_idx = out_v.max(dim=1)[1].item()\n",
        "            if np.random.random() < epsilon:\n",
        "                action_idx = env.action_space.sample()\n",
        "            action = Actions(action_idx)\n",
        "\n",
        "            close_price = env._state._cur_close()\n",
        "\n",
        "            if action == Actions.Buy and position is None:\n",
        "                position = close_price\n",
        "                position_steps = 0\n",
        "            elif action == Actions.Close and position is not None:\n",
        "                profit = close_price - position - (close_price + position) * comission / 100\n",
        "                profit = 100.0 * profit / position\n",
        "                stats['order_profits'].append(profit)\n",
        "                stats['order_steps'].append(position_steps)\n",
        "                position = None\n",
        "                position_steps = None\n",
        "\n",
        "            obs, reward, done, _ = env.step(action_idx)\n",
        "            total_reward += reward\n",
        "            episode_steps += 1\n",
        "            if position_steps is not None:\n",
        "                position_steps += 1\n",
        "            if done:\n",
        "                if position is not None:\n",
        "                    profit = close_price - position - (close_price + position) * comission / 100\n",
        "                    profit = 100.0 * profit / position\n",
        "                    stats['order_profits'].append(profit)\n",
        "                    stats['order_steps'].append(position_steps)\n",
        "                break\n",
        "\n",
        "        stats['episode_reward'].append(total_reward)\n",
        "        stats['episode_steps'].append(episode_steps)\n",
        "\n",
        "    return { key: np.mean(vals) for key, vals in stats.items() }"
      ],
      "metadata": {
        "id": "Y0K7eALDp3RD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Loaing Data for Trading Environment\n"
      ],
      "metadata": {
        "id": "X5_4V4z3oX9y"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "Prices = collections.namedtuple('Prices', field_names=['open', 'high', 'low', 'close', 'volume'])\n",
        "\n",
        "\n",
        "def read_csv(file_name, sep=',', filter_data=True, fix_open_prices=False):\n",
        "  print(\"Reading\", file_name)\n",
        "  with open(file_name, 'r') as fd:\n",
        "    reader = csv.reader(fd)\n",
        "    h = next(reader)\n",
        "    if 'Open' not in h and sep == ',':\n",
        "      return read_csv(file_name, ';')\n",
        "    indices = [h.index(s) for s in ('Open', 'High', 'Low', 'Close', 'Volume')]\n",
        "    #indices = [h.index(s) for s in ('<OPEN>', '<HIGH>', '<LOW>', '<CLOSE>', '<VOL>')]\n",
        "    o, h, l, c, v = [], [], [], [], []\n",
        "    count_out = 0\n",
        "    count_filter = 0 \n",
        "    count_fixed = 0\n",
        "    prev_vals = None\n",
        "    for row in reader:\n",
        "      vals = list(map(float, [row[idx] for idx in indices])) \n",
        "      if filter_data and all(map(lambda v: abs(v-vals[0]) < 1e-8, vals[:-1])):\n",
        "        count_filter += 1\n",
        "        continue\n",
        "      \n",
        "      po, ph, pl, pc, pv = vals\n",
        "      \n",
        "      # putting price data into list and then into a np.array \n",
        "      # where o is open price, c is close price, h is high price, l \n",
        "      # is low price, and v is volume\n",
        "      count_out +=1\n",
        "      o.append(po)\n",
        "      c.append(pc)\n",
        "      h.append(ph)\n",
        "      l.append(pl)\n",
        "      v.append(pv)\n",
        "      prev_vals = vals\n",
        "  return Prices(open=np.array(o, dtype=np.float32),high=np.array(h, dtype=np.float32), low=np.array(l, dtype=np.float32),close=np.array(c, dtype=np.float32), volume=np.array(v, dtype=np.float32))\n",
        "\n",
        "\n",
        "# prices(object): of collections.namedtuple type\n",
        "def prices_to_relative(prices):\n",
        "    \"\"\"\n",
        "    Convert prices to relative in respect to open price\n",
        "    :param ochl: tuple with open, close, high, low\n",
        "    :return: tuple with open, rel_close, rel_high, rel_low\n",
        "    \"\"\"\n",
        "    assert isinstance(prices, Prices)\n",
        "    rh = (prices.high - prices.open) / prices.open\n",
        "    rl = (prices.low - prices.open) / prices.open\n",
        "    rc = (prices.close - prices.open) / prices.open\n",
        "    return Prices(open=prices.open, high=rh, low=rl, close=rc, volume=prices.volume)\n",
        "\n",
        "def load_relative(csv_file):\n",
        "    return prices_to_relative(read_csv(csv_file))\n",
        "\n",
        "\n",
        "def price_files(dir_name):\n",
        "    result = []\n",
        "    for path in glob.glob(os.path.join(dir_name, \"*.csv\")):\n",
        "        result.append(path)\n",
        "    return result\n",
        "\n",
        "\n",
        "def load_year_data(year, basedir='data'):\n",
        "    y = str(year)[-2:]\n",
        "    result = {}\n",
        "    for path in glob.glob(os.path.join(basedir, \"*_%s*.csv\" % y)):\n",
        "        result[path] = load_relative(path)\n",
        "    return result\n"
      ],
      "metadata": {
        "id": "p5FOgwwMctR3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating Trading Environment using gym.Env \n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "gzAKILhh7lG_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# defualt number of past trading days agent can observe when taking action;\n",
        "# for 1D convolution model, this is the column portion of 2D matrix\n",
        "# default is 5 days\n",
        "BARS_COUNT = 5\n",
        "# default percentage of stock price trading agent pays broker when \n",
        "# buying/selling, default is 0.1% (i.e. very reasonable)\n",
        "DEFAULT_COMMISSION_PERC = 0.1\n",
        "\n",
        "\n",
        "# Actions\n",
        "class Actions(enum.Enum):\n",
        "  # actions agent can take when trading\n",
        "  Hold = 0 \n",
        "  Buy = 1\n",
        "  Close = 2\n",
        "\n",
        "# StocksEnv\n",
        "class StocksEnv(gym.Env):\n",
        "\n",
        "  # fields required by gym.Env\n",
        "  metadata = {'render.modes': ['human']}\n",
        "  spec = EnvSpec(\"SPYEnv-v0\")\n",
        "\n",
        "  def __init__(self, prices, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=True, state_1d=False, random_ofs_on_reset=True,\n",
        "               reward_on_close=False, volumes=True):\n",
        "    assert isinstance(prices, dict)\n",
        "    self._prices = prices\n",
        "\n",
        "#---------------------State-Observation Encoding Section------------------------    \n",
        "    # key!!!: creating the state observation for trading agent; when using\n",
        "    # the 1D convolutional model, encoding must be of State1D class!!!!\n",
        "    if state_1d:\n",
        "      self._state = State1D(bars_count, commission, reset_on_close, reward_on_close=reward_on_close, volumes=volumes)\n",
        "    else:\n",
        "      self._state = State(bars_count, commission, reset_on_close, reward_on_close=reward_on_close, volumes=volumes)\n",
        "    \n",
        "    # creating action space for trading agent\n",
        "    self.action_space = gym.spaces.Discrete(n=len(Actions))\n",
        "    \n",
        "    # creating observation space for training agent\n",
        "    self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=self._state.shape, dtype=np.float32)\n",
        "    \n",
        "    # decide if want to use random offset, default is True\n",
        "    self.random_ofs_on_reset = random_ofs_on_reset\n",
        "    self.seed()\n",
        "\n",
        "#------------------------Reset Section------------------------------------------\n",
        "  # creates the offset for time series data (i.e. not \n",
        "  # always starting at the beggining of the time series data for episode)\n",
        "  def reset(self):\n",
        "    self._instrument = self.np_random.choice(list(self._prices.keys()))\n",
        "    prices = self._prices[self._instrument]\n",
        "    bars = self._state.bars_count\n",
        "    if self.random_ofs_on_reset:\n",
        "      offset = self.np_random.choice(prices.high.shape[0]-bars*10)+bars\n",
        "    else:\n",
        "      offset = bars\n",
        "    self._state.reset(prices, offset)\n",
        "    return self._state.encode()  \n",
        "\n",
        "#-----------------------Step Section--------------------------------------------\n",
        "  # executes the sequence of agent taking action, getting reward and\n",
        "  # then getting the next observation/state \n",
        "  def step(self, action_idx):\n",
        "    action = Actions(action_idx)\n",
        "    reward, done = self._state.step(action)\n",
        "    obs = self._state.encode()\n",
        "    info = {\"instrument\":self._instrument, \"offset\": self._state._offset}\n",
        "    self._state.render(reward)\n",
        "    return obs, reward, done, info\n",
        "\n",
        "#----------------------Render Section-------------------------------------------\n",
        "  # required by gym.Env object; future will implement the render method to view\n",
        "  # the observation space of agent when trading to compare analysis\n",
        "  def render(self, mode='human', close=False):\n",
        "    pass  \n",
        "  def close(self):\n",
        "    pass\n",
        "\n",
        "  def seed(self, seed=None):\n",
        "    self.np_random, seed1 = seeding.np_random(seed)\n",
        "    seed2 = seeding.hash_seed(seed1+1) % 2**33\n",
        "    return [seed1, seed2]\n",
        "\n",
        "#----------------------Class Method Section-------------------------------------\n",
        "  # creates the instance of the  StocksEnv object (can be used with differnt\n",
        "  # stock data sets for multiple environment)\n",
        "  @classmethod\n",
        "  def from_dir(cls, data_dir, **kwargs):\n",
        "    prices = {f: load_relative(f) for f in price_files(data_dir)}\n",
        "    return StocksEnv(prices, **kwargs)"
      ],
      "metadata": {
        "id": "uDHY94JWfUMn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# State Space\n",
        "\n",
        "The reward is of ***either/or form***. This ***either/or form*** of the reward is done by setting the variable reward_on_close to either True or False, which \n",
        "is a Boolean parameter that switches between the two reward schemes. If it is set to True, the agent will receive a reward only on the close/selling of its stock position. If set to false, the agent will recive a reward only the buying and holding of its stock position (i.e. not selling).The default setting is True for reward_on_close which amounts to the trading strategy of [active investing](https://www.investopedia.com/terms/a/activeinvesting.asp#:~:text=Active%20investing%20refers%20to%20an,activity%20to%20exploit%20profitable%20conditions.). And Changing reward_on_close to False amounts to the trading strategy of [passive investing](https://www.investopedia.com/terms/p/passiveinvesting.asp#:~:text=Passive%20investing's%20goal%20is%20to,price%20fluctuations%20or%20market%20timing.)\n"
      ],
      "metadata": {
        "id": "Vei7smuIDjlo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# General State Class (i.e. models not based on convolutions)\n",
        "class State:\n",
        "  def __init__(self, bars_count, commission_perc, reset_on_close, reward_on_close=True, volumes=True):\n",
        "    super(State, self).__init__()\n",
        "    # checking bars_count is an int\n",
        "    assert isinstance(bars_count, int)\n",
        "    # checking that bars_count is greater than zero\n",
        "    assert bars_count > 0\n",
        "    # checking commission is a float\n",
        "    assert isinstance(commission_perc, float)\n",
        "    # checking commission is greater than zero\n",
        "    assert commission_perc >= 0.0\n",
        "    # checking that reset_on_close and reward on close are bools\n",
        "    assert isinstance(reset_on_close, bool)\n",
        "    assert isinstance(reward_on_close, bool)\n",
        "    self.bars_count=bars_count\n",
        "    self.commission_perc = commission_perc\n",
        "    self.reset_on_close = reset_on_close\n",
        "    self.reward_on_close = reward_on_close\n",
        "    self.volumes = volumes\n",
        "    self.bought_price = 0.0\n",
        "  \n",
        "  # method that reset's the environment \n",
        "  def reset(self, prices, offset):\n",
        "    assert isinstance(prices, Prices)\n",
        "    assert offset >= self.bars_count-1\n",
        "    self.have_position = False  # start with no stocks\n",
        "    self.open_price = 0.0\n",
        "    self._prices = prices\n",
        "    self._offset = offset\n",
        "\n",
        "  # shape of observation space is 1D\n",
        "  @property\n",
        "  def shape(self):\n",
        "    # the shape is the high, low, and closing prices of the current trading day\n",
        "    # (i.e. 3 or 4 if volume is used) times the num of bars\n",
        "    # (i.e. past prices agent can observe) plus the position flag \n",
        "    # (i.e. whether agent is holding onto the stock or not) and \n",
        "    # the relative profit agent has recieved since opening\n",
        "    if self.volumes:\n",
        "      return (5*self.bars_count+1+1, )\n",
        "    else:\n",
        "      return (4*self.bars_count+1+1, )\n",
        "  \n",
        "  # encoding observation space as a 1d Vector\n",
        "  def encode(self):\n",
        "    res = np.ndarray(shape=self.shape, dtype=np.float32)\n",
        "    shift = 0\n",
        "    for bar_idx in range(-self.bars_count+1, 1):\n",
        "      res[shift] = self._prices.open[self._offset + bar_idx]\n",
        "      shift += 1\n",
        "      res[shift] = self._prices.high[self._offset + bar_idx]\n",
        "      shift += 1\n",
        "      res[shift] = self._prices.low[self._offset + bar_idx]\n",
        "      shift += 1\n",
        "      res[shift] = self._prices.close[self._offset + bar_idx]\n",
        "      shift += 1\n",
        "      if self.volumes:\n",
        "        res[shift] = self._prices.volume[self._offset + bar_idx]\n",
        "        shift += 1\n",
        "    res[shift] = float(self.have_position)\n",
        "    shift += 1\n",
        "    if not self.have_position:\n",
        "      res[shift] = 0.0\n",
        "    else:\n",
        "      res[shift] = (self._cur_close() - self.open_price) / self.open_price\n",
        "    return res\n",
        " \n",
        "  def _cur_close(self):\n",
        "    \"\"\"\n",
        "    Calculate real close price for the current bar\n",
        "    \"\"\"\n",
        "    open = self._prices.open[self._offset]\n",
        "    rel_close = self._prices.close[self._offset]\n",
        "    return open * (1.0 + rel_close)\n",
        "\n",
        "#---------------!!!Step Section & Reward Calculation!!!------------------------- \n",
        "  def step(self, action):\n",
        "      \"\"\"\n",
        "      Perform one step in our price, adjust offset, check for the end of prices\n",
        "      and handle position change\n",
        "      :param action:\n",
        "      :return: reward, done\n",
        "      \"\"\"\n",
        "      assert isinstance(action, Actions)\n",
        "      reward = 0.0\n",
        "      done = False\n",
        "      close = self._cur_close()\n",
        "\n",
        "      if action == Actions.Buy and not self.have_position:\n",
        "        self.have_position = True\n",
        "        self.open_price = close\n",
        "        self.bought_price = close\n",
        "        reward -= self.commission_perc\n",
        "      elif action == Actions.Close and self.have_position:\n",
        "        reward -= self.commission_perc\n",
        "        done |= self.reset_on_close\n",
        "        \"\"\"\n",
        "        implements the active investing strategy,since a reward is only given \n",
        "        when selling;a positive reward is given when selling higher than when \n",
        "        you bought and a negative reward is given when selling lower than when \n",
        "        you bought the stock \n",
        "        \"\"\"\n",
        "        if self.reward_on_close:\n",
        "          reward += 100.0 * (close - self.open_price) / self.open_price\n",
        "        self.have_position = False\n",
        "        self.open_price = 0.0\n",
        "\n",
        "      self._offset += 1\n",
        "      prev_close = close\n",
        "      close = self._cur_close()\n",
        "      done |= self._offset >= self._prices.close.shape[0]-1\n",
        "\n",
        "      \"\"\"\n",
        "      implements the passive investing strategy ,since you only get this reward \n",
        "      when not selling; you get a positive reward when the current closing price \n",
        "      is higher than the price you bought the stock and you get a negative \n",
        "      reward when the current closing price is lower than the price you bought \n",
        "      the stock   \n",
        "      \"\"\"\n",
        "      if self.have_position and not self.reward_on_close:\n",
        "        reward += 100.0 * (close - self.bought_price) / self.bought_price\n",
        "\n",
        "      return reward, done\n",
        "  \n",
        "  # !!!!!TODO!!!rendering agent's observation space to screen!!!!!!!!!!!!!!!!!!\n",
        "  def render(self, reward, mode='human', close=False):\n",
        "    # creating plot object\n",
        "    # plt.clf()\n",
        "    # ofs = self.bars_count-1\n",
        "    # plt.plot(self._prices.open[self._offset-ofs:self._offset+1])\n",
        "    # plt.title(\"opening prices, data=%s\" % \"SPY ETF\")\n",
        "    # plt.ylabel(\"prices, %\")\n",
        "    # plt.show()\n",
        "    pass\n",
        "  \n",
        "  def close(self):\n",
        "    pass\n",
        "\n",
        "# Specific State Class for encoding observation space for convolution models\n",
        "class State1D(State):\n",
        "    \"\"\"\n",
        "    State with shape suitable for 1D convolution, must be 2D of form \n",
        "    (row, column) where row is either 6 for just price data and vol or \n",
        "    7 depending on position\n",
        "    \"\"\"\n",
        "\n",
        "    # shape of observation space is 2D\n",
        "    @property\n",
        "    def shape(self):\n",
        "        if self.volumes:\n",
        "            return (7, self.bars_count)\n",
        "        else:\n",
        "            return (6, self.bars_count)\n",
        "\n",
        "    # encoding observation space as a 2d Matrix\n",
        "    def encode(self):\n",
        "        res = np.zeros(shape=self.shape, dtype=np.float32)\n",
        "        ofs = self.bars_count-1\n",
        "        res[0] = self._prices.open[self._offset-ofs:self._offset+1]\n",
        "        res[1] = self._prices.high[self._offset-ofs:self._offset+1]\n",
        "        res[2] = self._prices.low[self._offset-ofs:self._offset+1]\n",
        "        res[3] = self._prices.close[self._offset-ofs:self._offset+1]\n",
        "        if self.volumes:\n",
        "            res[4] = self._prices.volume[self._offset-ofs:self._offset+1]\n",
        "            dst = 5\n",
        "        else:\n",
        "            dst = 4\n",
        "        if self.have_position:\n",
        "            res[dst] = 1.0\n",
        "            res[dst+1] = (self._cur_close() - self.open_price) / self.open_price\n",
        "        return res"
      ],
      "metadata": {
        "id": "_cAFNoOgbJ7P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating the Dueling DQN Network using 1D Convolutions \n"
      ],
      "metadata": {
        "id": "J2MAClP9fIV1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class DQNConv1D(nn.Module):\n",
        "    def __init__(self, shape, actions_n, kernel):\n",
        "        super(DQNConv1D, self).__init__()\n",
        "        \n",
        "        # 1D Convolutions for first transformation (i.e. cross-correlation)\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv1d(shape[0], 128, kernel),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv1d(128, 128, kernel),\n",
        "            nn.ReLU(),\n",
        "        )\n",
        "\n",
        "        out_size = self._get_conv_out(shape)\n",
        "\n",
        "        # linear transformation for feed-forward value \n",
        "        self.fc_val = nn.Sequential(\n",
        "            nn.Linear(out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 1)\n",
        "        )\n",
        "\n",
        "        # linear transformation for feed-forward advantage\n",
        "        self.fc_adv = nn.Sequential(\n",
        "            nn.Linear(out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, actions_n)\n",
        "        )\n",
        "\n",
        "    def _get_conv_out(self, shape):\n",
        "        o = self.conv(torch.zeros(1, *shape))\n",
        "        return int(np.prod(o.size()))\n",
        "\n",
        "    def forward(self, x):\n",
        "        conv_out = self.conv(x).view(x.size()[0], -1)\n",
        "        val = self.fc_val(conv_out)\n",
        "        adv = self.fc_adv(conv_out)\n",
        "        return val + adv - adv.mean(dim=1, keepdim=True)\n",
        "\n"
      ],
      "metadata": {
        "id": "PCJBg4JDmE43"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating the Policy Gradient Network using 1D Convolutions"
      ],
      "metadata": {
        "id": "DlGNDKVmNQys"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class PGN(nn.Module):\n",
        "    def __init__(self, shape, n_actions, kernel):\n",
        "        super(PGN, self).__init__()\n",
        "        \n",
        "        # 1D Convolutions for first transformation (i.e. cross-correlation)\n",
        "        self.conv_1 = nn.Sequential(\n",
        "            nn.Conv1d(shape[0], 128, kernel),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv1d(128, 128, kernel),\n",
        "            nn.ReLU()\n",
        "        )\n",
        "        \n",
        "        # linear transformation for policy net\n",
        "        self.policy = nn.Sequential(\n",
        "            nn.Linear(128, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, n_actions)\n",
        "        )\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        output = self.conv_1(x)\n",
        "        output = output.reshape(output.size(0), -1)\n",
        "        output = self.policy(output)\n",
        "        return output"
      ],
      "metadata": {
        "id": "EzyNfF2yNQKF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating A2C Network using 1D Convolutions "
      ],
      "metadata": {
        "id": "jbAHr9rIMTxA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class A2C(nn.Module):\n",
        "    def __init__(self, shape, actions_n, kernel):\n",
        "        super(A2C, self).__init__()\n",
        "\n",
        "        # 1D Convolutions for first transformation (i.e. cross-correlation)\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv1d(shape[0], 128, kernel),\n",
        "            nn.ReLU(),\n",
        "            nn.Conv1d(128, 128, kernel),\n",
        "            nn.ReLU(),\n",
        "        )\n",
        "        \n",
        "        conv_out_size = self._get_conv_out(shape)\n",
        "        \n",
        "        self.policy = nn.Sequential(\n",
        "            nn.Linear(conv_out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, actions_n)\n",
        "        )\n",
        "\n",
        "        self.value = nn.Sequential(\n",
        "            nn.Linear(conv_out_size, 512),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(512, 1)\n",
        "        )\n",
        "\n",
        "    def _get_conv_out(self, shape):\n",
        "        o = self.conv(torch.zeros(1, *shape))\n",
        "        return int(np.prod(o.size()))\n",
        "\n",
        "    def forward(self, x):\n",
        "        conv_out = self.conv(x).view(x.size()[0], -1)\n",
        "        return self.policy(conv_out), self.value(conv_out)"
      ],
      "metadata": {
        "id": "z56idH8gMS2Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating Baseline for Policy Gradient Method\n",
        "*   The Mean Buffer stores a moving average of 1 million transitions as the baseline which is subtracted from the discounted reward portion"
      ],
      "metadata": {
        "id": "SVwi_oMQibU_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MeanBuffer:\n",
        "    def __init__(self, capacity):\n",
        "        self.capacity = capacity\n",
        "        self.deque = collections.deque(maxlen=capacity)\n",
        "        self.sum = 0.0\n",
        "\n",
        "    def add(self, val):\n",
        "        if len(self.deque) == self.capacity:\n",
        "            self.sum -= self.deque[0]\n",
        "        self.deque.append(val)\n",
        "        self.sum += val\n",
        "\n",
        "    def mean(self):\n",
        "        if not self.deque:\n",
        "            return 0.0\n",
        "        return self.sum / len(self.deque)"
      ],
      "metadata": {
        "id": "aceaW5Myf2SI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating concurrent environmnets for Policy Gradient Method, A2C and A3C Method for correlated sampling problem "
      ],
      "metadata": {
        "id": "0CCaSZvhmEuH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiEnvs(gym.Wrapper):\n",
        "    def __init__(self, env):\n",
        "      super(MultiEnvs, self).__init__(env)\n",
        "    \n",
        "    def render(self, mode='human', close=True):\n",
        "      pass\n",
        "    \n",
        "    def close(self):\n",
        "      pass\n",
        "    \n",
        "def wrap_dqn(env):\n",
        "    env = MultiEnvs(env)\n",
        "    return env"
      ],
      "metadata": {
        "id": "lX4jd5ix8Dg9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def make_env(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=True, state_1d=True, random_ofs_on_reset=True,\n",
        "               reward_on_close=True, volumes=False, val=False):\n",
        "  if not val:\n",
        "    env = StocksEnv(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL)\n",
        "    env = gym.wrappers.TimeLimit(env, max_episode_steps=MAX_EPISODES)\n",
        "    return wrap_dqn(env)\n",
        "  else:\n",
        "    env = StocksEnv(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL)\n",
        "    env = gym.wrappers.TimeLimit(env, max_episode_steps=MAX_EPISODES)\n",
        "    return wrap_dqn(env)"
      ],
      "metadata": {
        "id": "a9IO3n8UmEXv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training the Trading Agent with the RewardTracker and TBMeanTracker Class(i.e. tracks agent status in environment)\n",
        "\n",
        "\n",
        "*   These classes track the trading agent while interacting in the environment. These were taken from the ptan common library and can be found here: [common](https://github.com/Shmuma/ptan/blob/master/ptan/common/utils.py)\n",
        "\n"
      ],
      "metadata": {
        "id": "eN4fLLywnXk5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class RewardTracker:\n",
        "\n",
        "    # stop_reward(int): reward stopping threshold for agent, defualt is 1 \n",
        "    # group_rewards(int): trading period, default is 1 trading day\n",
        "    def __init__(self, writer, stop_reward=1, group_rewards=1):\n",
        "        self.writer = writer\n",
        "        self.stop_reward = stop_reward\n",
        "        self.reward_buf = []\n",
        "        self.steps_buf = []\n",
        "        self.group_rewards = group_rewards\n",
        "\n",
        "    def __enter__(self):\n",
        "        self.ts = time.time()\n",
        "        self.ts_frame = 0\n",
        "        self.total_rewards = []\n",
        "        self.done_episodes = 0\n",
        "        self.total_steps = []\n",
        "        return self\n",
        "\n",
        "    def __exit__(self, *args):\n",
        "        self.writer.close()\n",
        "\n",
        "#-----------------------------Reward Section------------------------------------\n",
        "    def reward(self, reward_steps, frame, epsilon=None):\n",
        "        \n",
        "        # Reward Per Steps\n",
        "        reward, steps = reward_steps\n",
        "       \n",
        "        self.reward_buf.append(reward)\n",
        "        self.steps_buf.append(steps)\n",
        "        if len(self.reward_buf) < self.group_rewards:\n",
        "            return False\n",
        "        # calculates mean reward from buffer \n",
        "        reward = np.mean(self.reward_buf)\n",
        "        steps = np.mean(self.steps_buf)\n",
        "        self.reward_buf.clear()\n",
        "        self.steps_buf.clear()\n",
        "\n",
        "        # Total Rewards\n",
        "        self.total_rewards.append(reward)\n",
        "\n",
        "        # Total Steps\n",
        "        self.total_steps.append(steps)\n",
        "\n",
        "        speed = (frame - self.ts_frame) / (time.time() - self.ts)\n",
        "        self.ts_frame = frame\n",
        "        self.ts = time.time()\n",
        "\n",
        "        # Mean Reward \n",
        "        mean_reward = np.mean(self.total_rewards[-self.group_rewards:]) \n",
        "\n",
        "        mean_steps = np.mean(self.total_steps[-self.group_rewards:])\n",
        "        epsilon_str = \"\" if epsilon is None else \", eps %.2f\" % epsilon\n",
        "        print(\"%d: done %d games, mean reward %.3f, mean steps %.2f, speed %.2f f/s%s\" % (\n",
        "            frame, len(self.total_rewards)*self.group_rewards, mean_reward, mean_steps, speed, epsilon_str\n",
        "        ))\n",
        "        sys.stdout.flush()\n",
        "        if epsilon is not None:\n",
        "            self.writer.add_scalar(\"epsilon\", epsilon, frame)\n",
        "        self.writer.add_scalar(\"reward_per_100_tradingWindow\", mean_reward, frame)\n",
        "        self.writer.add_scalar(\"reward\", reward, frame)\n",
        "        self.writer.add_scalar(\"steps_per_100_tradingWindow\", mean_steps, frame)\n",
        "        if mean_reward > self.stop_reward:\n",
        "            print(\"Kid, you're on a roll. Enjoy it while it lasts, 'cause it never does.\")\n",
        "            return True\n",
        "        return False\n",
        "    \n",
        "    def reward_two(self, reward, frame, done_episodes):\n",
        "      # Reward Per Steps\n",
        "      self.total_rewards.append(reward)\n",
        "      mean_rewards = float(np.mean(self.total_rewards[-self.group_rewards:]))\n",
        "      print(\"%d: reward: %6.2f, mean_100: %6.2f, episodes: %d\" % (frame, reward, mean_rewards, done_episodes))\n",
        "      sys.stdout.flush()\n",
        "      self.writer.add_scalar(\"reward\", reward, frame)\n",
        "      self.writer.add_scalar(\"reward_100\", mean_rewards, frame)\n",
        "      self.writer.add_scalar(\"episodes\", done_episodes, frame)\n",
        "      if mean_rewards > self.stop_reward:\n",
        "        print(\"Kid, you're on a roll. Enjoy it while it lasts, 'cause it never does.\")\n",
        "        return True\n",
        "      return False\n",
        "\n",
        "    def reward_three(self, reward_steps, frame, epsilon=None):\n",
        "        \n",
        "        # Reward Per Steps\n",
        "        reward = reward_steps\n",
        "       \n",
        "        self.reward_buf.append(reward)\n",
        "        self.steps_buf.append(frame)\n",
        "        if len(self.reward_buf) < self.group_rewards:\n",
        "            return False\n",
        "        # calculates mean reward from buffer \n",
        "        reward = np.mean(self.reward_buf)\n",
        "        steps = np.mean(self.steps_buf)\n",
        "        self.reward_buf.clear()\n",
        "        self.steps_buf.clear()\n",
        "\n",
        "        # Total Rewards\n",
        "        self.total_rewards.append(reward)\n",
        "\n",
        "        # Total Steps\n",
        "        self.total_steps.append(steps)\n",
        "\n",
        "        speed = (frame - self.ts_frame) / (time.time() - self.ts)\n",
        "        self.ts_frame = frame\n",
        "        self.ts = time.time()\n",
        "\n",
        "        # Mean Reward \n",
        "        mean_reward = np.mean(self.total_rewards[-self.group_rewards:]) \n",
        "\n",
        "        mean_steps = np.mean(self.total_steps[-self.group_rewards:])\n",
        "        epsilon_str = \"\" if epsilon is None else \", eps %.2f\" % epsilon\n",
        "        print(\"Step %d: Episode %d: Mean reward %.3f\" % (\n",
        "            frame, len(self.total_rewards)*self.group_rewards, mean_reward))\n",
        "        sys.stdout.flush()\n",
        "        if epsilon is not None:\n",
        "            self.writer.add_scalar(\"epsilon\", epsilon, frame)\n",
        "        self.writer.add_scalar(\"reward_per_100_tradingWindow\", mean_reward, frame)\n",
        "        self.writer.add_scalar(\"reward\", reward, frame)\n",
        "        self.writer.add_scalar(\"steps_per_100_tradingWindow\", mean_steps, frame)\n",
        "        if mean_reward > self.stop_reward:\n",
        "            print(\"Kid, you're on a roll. Enjoy it while it lasts, 'cause it never does.\")\n",
        "            return True\n",
        "        return False\n",
        "\n",
        "\n",
        "class TBMeanTracker:\n",
        "    \"\"\"\n",
        "    TensorBoard value tracker: allows to batch fixed amount of historical values and write their mean into TB\n",
        "    Designed and tested with pytorch-tensorboard in mind\n",
        "    \"\"\"\n",
        "    def __init__(self, writer, batch_size):\n",
        "        \"\"\"\n",
        "        :param writer: writer with close() and add_scalar() methods\n",
        "        :param batch_size: integer size of batch to track\n",
        "        \"\"\"\n",
        "        assert isinstance(batch_size, int)\n",
        "        assert writer is not None\n",
        "        self.writer = writer\n",
        "        self.batch_size = batch_size\n",
        "\n",
        "    def __enter__(self):\n",
        "        self._batches = collections.defaultdict(list)\n",
        "        return self\n",
        "\n",
        "    def __exit__(self, exc_type, exc_val, exc_tb):\n",
        "        self.writer.close()\n",
        "\n",
        "    @staticmethod\n",
        "    def _as_float(value):\n",
        "        assert isinstance(value, (float, int, np.ndarray, np.generic, torch.autograd.Variable)) or torch.is_tensor(value)\n",
        "        tensor_val = None\n",
        "        if isinstance(value, torch.autograd.Variable):\n",
        "            tensor_val = value.data\n",
        "        elif torch.is_tensor(value):\n",
        "            tensor_val = value\n",
        "\n",
        "        if tensor_val is not None:\n",
        "            return tensor_val.float().mean().item()\n",
        "        elif isinstance(value, np.ndarray):\n",
        "            return float(np.mean(value))\n",
        "        else:\n",
        "            return float(value)\n",
        "\n",
        "    def track(self, param_name, value, iter_index):\n",
        "        assert isinstance(param_name, str)\n",
        "        assert isinstance(iter_index, int)\n",
        "\n",
        "        data = self._batches[param_name]\n",
        "        data.append(self._as_float(value))\n",
        "\n",
        "        if len(data) >= self.batch_size:\n",
        "            self.writer.add_scalar(param_name, np.mean(data), iter_index)\n",
        "            data.clear()\n"
      ],
      "metadata": {
        "id": "4AkcT1EdL2sL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#----------------------Parameters Section---------------------------------------\n",
        "# transitions to store in buffer for PGN baseline \n",
        "BASELINE_STEPS = 1000000\n",
        "# num of observations for gradient update \n",
        "BATCH_SIZE = 32\n",
        "BATCH_SIZE_PGN = 32\n",
        "BATCH_SIZE_A2C = 32\n",
        "# when to sync weights of target network for DQN\n",
        "TARGET_NET_SYNC = 1000\n",
        "# discount factor\n",
        "GAMMA = 0.99\n",
        "# determines magnitude (i.e. norm) of gradient vector\n",
        "# important for varience reduction and policy changing too fast\n",
        "GRAD_L2_CLIP = 0.1\n",
        "# determines how many steps ahead to approximate of the total discounted \n",
        "# reward for every action\n",
        "REWARD_STEPS = 2\n",
        "REWARD_STEPS_PGN = 2\n",
        "REWARD_STEPS_A2C = 2\n",
        "# replay buffer size for DQN\n",
        "REPLAY_SIZE = 100000\n",
        "# intial num of samples for DQN replay buffer\n",
        "REPLAY_INITIAL = 10000\n",
        "# default learning rate for Adam optimizer\n",
        "LEARNING_RATE = 0.001\n",
        "# weight of beta parmater for entropy loss for PGN and A2C \n",
        "ENTROPY_BETA = .022\n",
        "EPSILON_START = 1.0\n",
        "EPSILON_STOP = 0.1\n",
        "EPSILON_STEPS = 1000000\n",
        "# how many concurrent envs to create\n",
        "ENV_COUNT = 32\n",
        "LAMBDA = 10\n",
        "# when to save progress of training agent\n",
        "CHECKPOINT_EVERY_STEP = 100000\n",
        "# when to perform validation\n",
        "VALIDATION_EVERY_STEP = 100000\n",
        "CUDA = True\n",
        "# date-time features, default is None\n",
        "YEAR = None\n",
        "STATES_TO_EVALUATE = 1000\n",
        "# controls saving of model parameters\n",
        "EVAL_EVERY_STEP = 1000\n",
        "# max time episode default is 1000 episodes\n",
        "MAX_EPISODES = 1000\n",
        "# selects model type in str form can be either\n",
        "# DQN, PGN, or A2C, default is A2C\n",
        "MODEL = \"A2C\"\n",
        "# trading window, default is 100 days\n",
        "WINDOW = 100\n",
        "# reward to hit, default is infinity\n",
        "STOP_REWARD = np.inf\n",
        "# volume features, default is False \n",
        "VOL = False\n",
        "# reset on close, default is True\n",
        "RESET = True\n",
        "# convolution based model, default is True\n",
        "STATE1D = True\n",
        "# random offset on reset, defualt is True\n",
        "RANDOM = True\n",
        "# reward on close, default is True (i.e. active investing)\n",
        "REW_CLOSE = True\n",
        "# kernel size, default is 3x3\n",
        "KERNEL = 3\n",
        "\n",
        "\n",
        "#----------------------Load Data Section----------------------------------------\n",
        "\n",
        "# dir/file of market data for training in csv form (i.e. past historical data) \n",
        "DEFAULT_STOCKS = \"/content/drive/MyDrive/Datasets/SPY/spy_past.csv\"\n",
        "\n",
        "# dir/file of market data for validation in csv form (i.e. past historical data) \n",
        "DEFAULT_VAL_STOCKS = \"/content/drive/MyDrive/Datasets/SPY/spy_future.csv\"\n",
        "SAVE_PATH = \"saves\"\n",
        "\n",
        "#-------DQN Calculation Section(i.e. loss and state calculation)----------------\n",
        "def calc_values_of_states(states, net, device=\"cpu\"):\n",
        "    mean_vals = []\n",
        "    for batch in np.array_split(states, 64):\n",
        "        states_v = torch.tensor(batch).to(device)\n",
        "        action_values_v = net(states_v)\n",
        "        best_action_values_v = action_values_v.max(1)[0]\n",
        "        mean_vals.append(best_action_values_v.mean().item())\n",
        "    return np.mean(mean_vals)\n",
        "\n",
        "\n",
        "def unpack_batch_one(batch):\n",
        "    states, actions, rewards, dones, last_states = [], [], [], [], []\n",
        "    for exp in batch:\n",
        "        state = np.array(exp.state, copy=False)\n",
        "        states.append(state)\n",
        "        actions.append(exp.action)\n",
        "        rewards.append(exp.reward)\n",
        "        dones.append(exp.last_state is None)\n",
        "        if exp.last_state is None:\n",
        "            last_states.append(state)       # the result will be masked anyway\n",
        "        else:\n",
        "            last_states.append(np.array(exp.last_state, copy=False))\n",
        "    return np.array(states, copy=False), np.array(actions), np.array(rewards, dtype=np.float32), \\\n",
        "           np.array(dones, dtype=np.uint8), np.array(last_states, copy=False)\n",
        "\n",
        "\n",
        "def calc_loss(batch, net, tgt_net, gamma, device=\"cpu\"):\n",
        "    states, actions, rewards, dones, next_states = unpack_batch_one(batch)\n",
        "\n",
        "    states_v = torch.tensor(states).to(device)\n",
        "    next_states_v = torch.tensor(next_states).to(device)\n",
        "    actions_v = torch.tensor(actions).to(device)\n",
        "    rewards_v = torch.tensor(rewards).to(device)\n",
        "    done_mask = torch.ByteTensor(dones).to(device)\n",
        "\n",
        "    state_action_values = net(states_v).gather(1, actions_v.unsqueeze(-1)).squeeze(-1)\n",
        "    next_state_actions = net(next_states_v).max(1)[1]\n",
        "    next_state_values = tgt_net(next_states_v).gather(1, next_state_actions.unsqueeze(-1)).squeeze(-1)\n",
        "    next_state_values[done_mask] = 0.0\n",
        "\n",
        "    expected_state_action_values = next_state_values.detach() * gamma + rewards_v\n",
        "    return nn.MSELoss()(state_action_values, expected_state_action_values)\n",
        "\n",
        "#---------------------A2C Calculation Section-----------------------------------\n",
        "def unpack_batch(batch, net, device='cpu'):\n",
        "    \"\"\"\n",
        "    Convert batch into training tensors\n",
        "    :param batch:\n",
        "    :param net:\n",
        "    :return: states variable, actions tensor, reference values variable\n",
        "    \"\"\"\n",
        "    states = []\n",
        "    actions = []\n",
        "    rewards = []\n",
        "    not_done_idx = []\n",
        "    last_states = []\n",
        "    for idx, exp in enumerate(batch):\n",
        "        states.append(np.array(exp.state, copy=False))\n",
        "        actions.append(int(exp.action))\n",
        "        rewards.append(exp.reward)\n",
        "        if exp.last_state is not None:\n",
        "            not_done_idx.append(idx)\n",
        "            last_states.append(np.array(exp.last_state, copy=False))\n",
        "    states_v = torch.FloatTensor(np.array(states, copy=False)).to(device)\n",
        "    actions_t = torch.LongTensor(actions).to(device)\n",
        "    # handle rewards\n",
        "    rewards_np = np.array(rewards, dtype=np.float32)\n",
        "    if not_done_idx:\n",
        "        last_states_v = torch.FloatTensor(np.array(last_states, copy=False)).to(device)\n",
        "        last_vals_v = net(last_states_v)[1]\n",
        "        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]\n",
        "        rewards_np[not_done_idx] += GAMMA ** REWARD_STEPS_A2C  * last_vals_np\n",
        "\n",
        "    ref_vals_v = torch.FloatTensor(rewards_np).to(device)\n",
        "    return states_v, actions_t, ref_vals_v\n",
        "\n",
        "\n",
        "\n",
        "#-----------------------------Main Section--------------------------------------\n",
        "if __name__ == \"__main__\":\n",
        "    device = torch.device(\"cuda\" if CUDA else \"cpu\")\n",
        "    saves_path = os.path.join(\"/content/\", SAVE_PATH)\n",
        "    os.makedirs(saves_path, exist_ok=True)\n",
        "\n",
        "    # data load\n",
        "    if YEAR is not None or os.path.isfile(DEFAULT_STOCKS):\n",
        "        if YEAR is not None:\n",
        "            stock_data = data.load_year_data(YEAR)\n",
        "        else:\n",
        "            stock_data = {\"SPY\": load_relative(DEFAULT_STOCKS)}\n",
        "        \n",
        "        if MODEL==\"DQN\":\n",
        "          # training env\n",
        "          env = StocksEnv(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL)\n",
        "          # test env\n",
        "          env_tst = StocksEnv(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL)\n",
        "           # time limit \n",
        "          env = gym.wrappers.TimeLimit(env, max_episode_steps=MAX_EPISODES)\n",
        "          # validation env\n",
        "          val_data = {\"SPY\": load_relative(DEFAULT_VAL_STOCKS)}\n",
        "          env_val = StocksEnv(val_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL)\n",
        "        elif MODEL==\"PGN\":\n",
        "          # training envs\n",
        "          envs = [make_env(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL, val=False) for _ in range(ENV_COUNT)]\n",
        "          # test envs\n",
        "          envs_tst = [make_env(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL, val=False) for _ in range(ENV_COUNT)]\n",
        "          # validation envs\n",
        "          val_data = {\"SPY\": load_relative(DEFAULT_VAL_STOCKS)}\n",
        "          envs_vals = [make_env(val_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL, val=True) for _ in range(ENV_COUNT)]\n",
        "        elif MODEL==\"A2C\":\n",
        "           # training env\n",
        "          envs = [make_env(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL, val=False) for _ in range(ENV_COUNT)]\n",
        "           # test envs\n",
        "          envs_tst = [make_env(stock_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL, val=False) for _ in range(ENV_COUNT)]\n",
        "          # validation envs\n",
        "          val_data = {\"SPY\": load_relative(DEFAULT_VAL_STOCKS)}\n",
        "          envs_vals = [make_env(val_data, bars_count=BARS_COUNT, commission=DEFAULT_COMMISSION_PERC,\n",
        "               reset_on_close=RESET, state_1d=STATE1D, random_ofs_on_reset=RANDOM,\n",
        "               reward_on_close=REW_CLOSE, volumes=VOL, val=True) for _ in range(ENV_COUNT)]\n",
        "        else:\n",
        "          raise RuntimeError(\"Model does not exist\")\n",
        "    else:\n",
        "        raise RuntimeError(\"See Load Data Section above to input data for agent to train on\")\n",
        "    \n",
        "  \n",
        "    \n",
        "    # tensorboard stuff\n",
        "    writer = SummaryWriter(comment=\"-simple-\" + \"run\")\n",
        "\n",
        "    # model selection\n",
        "    if MODEL==\"DQN\":\n",
        "      net = DQNConv1D(env.observation_space.shape, env.action_space.n, KERNEL).to(device)\n",
        "      # optimizer\n",
        "      optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE) \n",
        "    elif MODEL==\"PGN\":\n",
        "      net =  PGN(envs[0].observation_space.shape, envs[0].action_space.n, KERNEL).to(device)\n",
        "      # optimizer\n",
        "      optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3) \n",
        "    elif MODEL==\"A2C\":\n",
        "      net =  A2C(envs[0].observation_space.shape, envs[0].action_space.n, KERNEL).to(device)\n",
        "      optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)\n",
        "\n",
        "#----------------------------Ptan Section---------------------------------------    \n",
        "    # https://github.com/Shmuma/ptan/blob/master/ptan/agent.py\n",
        "    tgt_net = ptan.agent.TargetNet(net)\n",
        "    # https://github.com/Shmuma/ptan/blob/master/ptan/actions.py\n",
        "    selector = ptan.actions.EpsilonGreedyActionSelector(EPSILON_START)\n",
        "    # https://github.com/Shmuma/ptan/blob/master/ptan/agent.py\n",
        "    if MODEL==\"DQN\":\n",
        "      agent = ptan.agent.DQNAgent(net, selector, device=device)\n",
        "      exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, GAMMA, steps_count=REWARD_STEPS)\n",
        "    elif MODEL==\"PGN\":\n",
        "      agent = ptan.agent.PolicyAgent(net,device=device,apply_softmax=True, preprocessor=ptan.agent.float32_preprocessor)\n",
        "      exp_source = ptan.experience.ExperienceSourceFirstLast(envs, agent, GAMMA, steps_count=REWARD_STEPS_PGN)\n",
        "    elif MODEL==\"A2C\":\n",
        "      agent = ptan.agent.ActorCriticAgent(net,device=device,apply_softmax=True)\n",
        "      exp_source = ptan.experience.ExperienceSourceFirstLast(envs, agent, GAMMA, steps_count=REWARD_STEPS_A2C)\n",
        "    # https://github.com/Shmuma/ptan/blob/master/ptan/experience.py\n",
        "    buffer = ptan.experience.ExperienceReplayBuffer(exp_source, REPLAY_SIZE)\n",
        "\n",
        "#-----------------------RewardTracker/Training Loop Section---------------------\n",
        "    # initialization\n",
        "    step_idx = 0\n",
        "    eval_states = None\n",
        "    best_mean_val = None\n",
        "    done_episodes = 0\n",
        "    reward_sum = 0.0\n",
        "    step_rewards = []\n",
        "    baseline_buf = MeanBuffer(BASELINE_STEPS)\n",
        "\n",
        "    batch_states, batch_actions, batch_scales = [], [], []\n",
        "\n",
        "    with RewardTracker(writer, stop_reward=STOP_REWARD, group_rewards=WINDOW) as reward_tracker:\n",
        "      # DQN training loop\n",
        "      if MODEL==\"DQN\":\n",
        "        while True:\n",
        "          step_idx += 1\n",
        "          buffer.populate(1)\n",
        "          selector.epsilon = max(EPSILON_STOP, EPSILON_START - step_idx / EPSILON_STEPS)\n",
        "          \n",
        "          new_rewards = exp_source.pop_rewards_steps()\n",
        "\n",
        "          if new_rewards:\n",
        "            reward_tracker.reward(new_rewards[0], step_idx, selector.epsilon)\n",
        "          \n",
        "          if len(buffer) < REPLAY_INITIAL:\n",
        "            continue\n",
        "          \n",
        "          # Replay Buffer\n",
        "          if eval_states is None:\n",
        "            print(\"Initial buffer populated, start training\")\n",
        "            eval_states = buffer.sample(STATES_TO_EVALUATE)\n",
        "            eval_states = [np.array(transition.state, copy=False) for transition in eval_states]\n",
        "            eval_states = np.array(eval_states, copy=False)\n",
        "            \n",
        "          # Training \n",
        "          if step_idx % EVAL_EVERY_STEP == 0:\n",
        "            mean_val = calc_values_of_states(eval_states, net, device=device)\n",
        "            writer.add_scalar(\"values_mean\", mean_val, step_idx)\n",
        "            if best_mean_val is None or best_mean_val < mean_val:\n",
        "              if best_mean_val is not None:\n",
        "                print(\"%d: Best mean value updated %.3f -> %.3f\" % (step_idx, best_mean_val, mean_val))\n",
        "              best_mean_val = mean_val\n",
        "              torch.save(net.state_dict(), os.path.join(saves_path, \"mean_val-%.3f.data\" % mean_val))\n",
        "              \n",
        "              if step_idx >= 4000000:\n",
        "                plt.clf()\n",
        "                plt.plot(rewards)\n",
        "                plt.title(\"Investing Total reward, data=%s\" % \"SPY ETF\")\n",
        "                plt.ylabel(\"Total Reward, %\")\n",
        "                plt.savefig(\"rewards-%s.png\" % \"passive investing\")\n",
        "                break\n",
        "\n",
        "          optimizer.zero_grad()\n",
        "          batch = buffer.sample(BATCH_SIZE)\n",
        "          loss_v = calc_loss(batch, net, tgt_net.target_model, GAMMA ** REWARD_STEPS, device=device)\n",
        "          loss_v.backward()\n",
        "          optimizer.step()\n",
        "\n",
        "          # Updating Target Weights\n",
        "          if step_idx % TARGET_NET_SYNC == 0:\n",
        "            tgt_net.sync()\n",
        "          \n",
        "          # Checkpointing\n",
        "          if step_idx % CHECKPOINT_EVERY_STEP == 0:\n",
        "            idx = step_idx // CHECKPOINT_EVERY_STEP\n",
        "            torch.save(net.state_dict(), os.path.join(saves_path, \"checkpoint-%3d.data\" % idx))\n",
        "            \n",
        "          # Validation and Testing\n",
        "          if step_idx % VALIDATION_EVERY_STEP == 0:\n",
        "            res = validation_run(env_tst, net, device=device)\n",
        "            for key, val in res.items():\n",
        "              writer.add_scalar(key + \"_test\", val, step_idx)\n",
        "            res = validation_run(env_val, net, device=device)\n",
        "            for key, val in res.items():\n",
        "              writer.add_scalar(key + \"_val\", val, step_idx)\n",
        "      # PGN training loop\n",
        "      elif MODEL==\"PGN\":\n",
        "        for step_idx, exp in enumerate(exp_source):\n",
        "          baseline_buf.add(exp.reward)\n",
        "          baseline = baseline_buf.mean()\n",
        "          batch_states.append(np.array(exp.state, copy=False))\n",
        "          batch_actions.append(int(exp.action))\n",
        "          batch_scales.append(exp.reward - baseline)\n",
        "          \n",
        "          \n",
        "          # handle new rewards\n",
        "          new_rewards = exp_source.pop_total_rewards()\n",
        "          if new_rewards:\n",
        "            done_episodes += 1\n",
        "            reward = new_rewards[0]\n",
        "            reward_tracker.reward_two(reward, step_idx, done_episodes)\n",
        "\n",
        "          if len(batch_states) < BATCH_SIZE_PGN:\n",
        "            continue\n",
        "          \n",
        "          states_v = torch.tensor(batch_states).to(device)\n",
        "          batch_actions_t = torch.tensor(batch_actions).to(device)\n",
        "\n",
        "          scale_std = np.std(batch_scales)\n",
        "          batch_scale_v = torch.tensor(batch_scales).to(device)\n",
        "          \n",
        "          # policy loss\n",
        "          optimizer.zero_grad()\n",
        "          logits_v = net(states_v)\n",
        "          log_prob_v = F.log_softmax(logits_v, dim=1)\n",
        "          log_prob_actions_v = batch_scale_v * log_prob_v[range(BATCH_SIZE_PGN), batch_actions_t]\n",
        "          loss_policy_v = -log_prob_actions_v.mean()\n",
        "\n",
        "          # entropy loss\n",
        "          prob_v = F.softmax(logits_v, dim=1)\n",
        "          entropy_v = -(prob_v * log_prob_v).sum(dim=1).mean()\n",
        "          entropy_loss_v = -ENTROPY_BETA * entropy_v\n",
        "\n",
        "          # calculate total loss\n",
        "          loss_v = loss_policy_v + entropy_loss_v \n",
        "          loss_v.backward()\n",
        "          nn_utils.clip_grad_norm_(net.parameters(), GRAD_L2_CLIP)\n",
        "          optimizer.step()\n",
        "\n",
        "          batch_states.clear()\n",
        "          batch_actions.clear()\n",
        "          batch_scales.clear()\n",
        "\n",
        "          # checkpointing\n",
        "          if step_idx % CHECKPOINT_EVERY_STEP == 0:\n",
        "            idx = step_idx // CHECKPOINT_EVERY_STEP\n",
        "            torch.save(net.state_dict(), os.path.join(saves_path, \"checkpoint-%3d.data\" % idx))\n",
        "            \n",
        "          # validation and testing\n",
        "          if step_idx % VALIDATION_EVERY_STEP == 0:\n",
        "            res = validation_run(envs_tst, net, device=device)\n",
        "            for key, val in res.items():\n",
        "              writer.add_scalar(key + \"_test\", val, step_idx)\n",
        "            res = validation_run(envs_vals, net, device=device)\n",
        "            for key, val in res.items():\n",
        "              writer.add_scalar(key + \"_val\", val, step_idx)\n",
        "      # A2C training loop\n",
        "      elif MODEL==\"A2C\":\n",
        "        with TBMeanTracker(writer, batch_size=10) as tb_tracker:\n",
        "          batch = []\n",
        "          for step_idx, exp in enumerate(exp_source):\n",
        "            \n",
        "            batch.append(exp)\n",
        "            # handle new rewards\n",
        "            new_rewards = exp_source.pop_total_rewards()\n",
        "            if new_rewards:\n",
        "              reward_tracker.reward_three(new_rewards[0], step_idx)\n",
        "\n",
        "            if len(batch) < BATCH_SIZE_A2C:\n",
        "              continue\n",
        "          \n",
        "            states_v, actions_t, vals_ref_v = unpack_batch(batch, net, device=device)\n",
        "            batch.clear()\n",
        " \n",
        "            # zero-ing gradients\n",
        "            optimizer.zero_grad()\n",
        "          \n",
        "            # critic/value loss \n",
        "            logits_v, value_v = net(states_v)\n",
        "            loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)\n",
        "          \n",
        "            # actor/policy loss\n",
        "            log_prob_v = F.log_softmax(logits_v, dim=1)\n",
        "            adv_v = vals_ref_v - value_v.squeeze(-1).detach()\n",
        "            log_prob_actions_v = adv_v * log_prob_v[range(BATCH_SIZE_A2C), actions_t]\n",
        "            loss_policy_v = -log_prob_actions_v.mean()\n",
        "\n",
        "            # entropy loss\n",
        "            prob_v = F.softmax(logits_v, dim=1)\n",
        "            entropy_v = -(prob_v * log_prob_v).sum(dim=1).mean()\n",
        "            entropy_loss_v = -ENTROPY_BETA * entropy_v\n",
        "          \n",
        "            # calculate policy gradient loss only\n",
        "            loss_policy_v.backward(retain_graph=True)\n",
        "            grads = np.concatenate([p.grad.data.cpu().numpy().flatten()\n",
        "                                        for p in net.parameters()\n",
        "                                        if p.grad is not None])\n",
        "\n",
        "            # calculate entropy and value gradients\n",
        "            loss_v = entropy_loss_v + loss_value_v\n",
        "            loss_v.backward()\n",
        "            nn_utils.clip_grad_norm_(net.parameters(), GRAD_L2_CLIP)\n",
        "            optimizer.step()\n",
        "          \n",
        "            # get full loss\n",
        "            loss_v += loss_policy_v\n",
        "\n",
        "            # checkpointing\n",
        "            if step_idx % CHECKPOINT_EVERY_STEP == 0:\n",
        "              idx = step_idx // CHECKPOINT_EVERY_STEP\n",
        "              torch.save(net.state_dict(), os.path.join(saves_path, \"checkpoint-%3d.data\" % idx))\n",
        "            \n",
        "            # validation and testing\n",
        "            if step_idx % VALIDATION_EVERY_STEP == 0:\n",
        "              res = validation_run(envs_tst, net, device=device)\n",
        "              for key, val in res.items():\n",
        "                writer.add_scalar(key + \"_test\", val, step_idx)\n",
        "              res = validation_run(envs_vals, net, device=device)\n",
        "              for key, val in res.items():\n",
        "                writer.add_scalar(key + \"_val\", val, step_idx)\n",
        "\n",
        "\n",
        "            tb_tracker.track(\"advantage\",       adv_v, step_idx)\n",
        "            tb_tracker.track(\"values\",          value_v, step_idx)\n",
        "            tb_tracker.track(\"batch_rewards\",   vals_ref_v, step_idx)\n",
        "            tb_tracker.track(\"loss_entropy\",    entropy_loss_v, step_idx)\n",
        "            tb_tracker.track(\"loss_policy\",     loss_policy_v, step_idx)\n",
        "            tb_tracker.track(\"loss_value\",      loss_value_v, step_idx)\n",
        "            tb_tracker.track(\"loss_total\",      loss_v, step_idx)\n",
        "            tb_tracker.track(\"grad_l2\",         np.sqrt(np.mean(np.square(grads))), step_idx)\n",
        "            tb_tracker.track(\"grad_max\",        np.max(np.abs(grads)), step_idx)\n",
        "            tb_tracker.track(\"grad_var\",        np.var(grads), step_idx)\n",
        "\n",
        "      else:\n",
        "        raise RuntimeError(\"Something went wrong, fix it!!!\")\n"
      ],
      "metadata": {
        "id": "VvCRry1ooHfz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir runs/"
      ],
      "metadata": {
        "id": "h7cdLOBVyV8l"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!zip -r /content/runs.zip /content/runs/"
      ],
      "metadata": {
        "id": "1nlZzU4M3o8t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# !tar -xvzf  /content/drive/MyDrive/Datasets/StockMarketData/ch08-small-quotes.tgz -C /content/"
      ],
      "metadata": {
        "id": "XJoiq84jCM3V"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}