test2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "mount_file_id": "1mydiVo6yhvO5w1ARwF2KVsLzOeqZ8wM4",
      "authorship_tag": "ABX9TyNrUktch0AZOdUbv7VNefRm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/shelan-de-livera/finalprojecttestrepo/blob/main/test2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install tf-agents"
      ],
      "metadata": {
        "id": "cex_P0kK5I6J",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ac859168-a041-4532-9182-04623708d59f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting tf-agents\n",
            "  Downloading tf_agents-0.18.0-py3-none-any.whl (1.4 MB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/1.4 MB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━\u001b[0m\u001b[90m╺\u001b[0m\u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.1/1.4 MB\u001b[0m \u001b[31m2.5 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m \u001b[32m1.4/1.4 MB\u001b[0m \u001b[31m20.7 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.4/1.4 MB\u001b[0m \u001b[31m17.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: absl-py>=0.6.1 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (1.4.0)\n",
            "Requirement already satisfied: cloudpickle>=1.3 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (2.2.1)\n",
            "Requirement already satisfied: gin-config>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (0.5.0)\n",
            "Collecting gym<=0.23.0,>=0.17.0 (from tf-agents)\n",
            "  Downloading gym-0.23.0.tar.gz (624 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m624.4/624.4 kB\u001b[0m \u001b[31m38.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: numpy>=1.19.0 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (1.23.5)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.10/dist-packages (from tf-agents) (9.4.0)\n",
            "Requirement already satisfied: six>=1.10.0 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (1.16.0)\n",
            "Requirement already satisfied: protobuf>=3.11.3 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (3.20.3)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (1.14.1)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (4.5.0)\n",
            "Collecting pygame==2.1.3 (from tf-agents)\n",
            "  Downloading pygame-2.1.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (13.7 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.7/13.7 MB\u001b[0m \u001b[31m36.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: tensorflow-probability~=0.22.0 in /usr/local/lib/python3.10/dist-packages (from tf-agents) (0.22.0)\n",
            "Requirement already satisfied: gym-notices>=0.0.4 in /usr/local/lib/python3.10/dist-packages (from gym<=0.23.0,>=0.17.0->tf-agents) (0.0.8)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability~=0.22.0->tf-agents) (4.4.2)\n",
            "Requirement already satisfied: gast>=0.3.2 in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability~=0.22.0->tf-agents) (0.5.4)\n",
            "Requirement already satisfied: dm-tree in /usr/local/lib/python3.10/dist-packages (from tensorflow-probability~=0.22.0->tf-agents) (0.1.8)\n",
            "Building wheels for collected packages: gym\n",
            "  Building wheel for gym (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for gym: filename=gym-0.23.0-py3-none-any.whl size=697629 sha256=0c9f154d59e9fdfd4b8a88fb77b1347207b9dc677014252fdb5c0bf4fadf6b7d\n",
            "  Stored in directory: /root/.cache/pip/wheels/3d/6f/b4/3991d4fae11d0ecb0754c11cc1b4e7745012850da4efaaf0b1\n",
            "Successfully built gym\n",
            "Installing collected packages: pygame, gym, tf-agents\n",
            "  Attempting uninstall: pygame\n",
            "    Found existing installation: pygame 2.5.2\n",
            "    Uninstalling pygame-2.5.2:\n",
            "      Successfully uninstalled pygame-2.5.2\n",
            "  Attempting uninstall: gym\n",
            "    Found existing installation: gym 0.25.2\n",
            "    Uninstalling gym-0.25.2:\n",
            "      Successfully uninstalled gym-0.25.2\n",
            "Successfully installed gym-0.23.0 pygame-2.1.3 tf-agents-0.18.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function"
      ],
      "metadata": {
        "id": "DRMn-MKv9UVR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import gin\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import\n",
        "import tensorflow_probability as tfp"
      ],
      "metadata": {
        "id": "dCZKbAsd9WOx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tf_agents.networks import categorical_projection_network\n",
        "from tf_agents.networks import network\n",
        "from tf_agents.networks import normal_projection_network\n",
        "from tf_agents.specs import tensor_spec\n",
        "from tf_agents.utils import nest_utils\n",
        "from tf_agents.networks import encoding_network\n",
        "from tf_agents.keras_layers import dynamic_unroll_layer\n",
        "from tf_agents.trajectories import time_step"
      ],
      "metadata": {
        "id": "0yB6quR59YrJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# agents - agents.agent.py\n"
      ],
      "metadata": {
        "id": "iS2SjIR0A26j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from abc import ABC, abstractmethod"
      ],
      "metadata": {
        "id": "tjtNxNaRAx4c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# from abc import ABC, abstractmethod\n",
        "\n",
        "class Agent(ABC):\n",
        "    @abstractmethod\n",
        "    def initialize(self):\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def save(self):\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def load(self):\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def train(self, **kwargs) -> list:\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def eval(self, **kwargs) -> float:\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def compute_action(self, **kwargs):\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def get_action_probabilities(self, **kwargs):\n",
        "        pass\n"
      ],
      "metadata": {
        "id": "AebQbstm_sNG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# agents.tfagents - tfagents.tfagent.py\n"
      ],
      "metadata": {
        "id": "UomrQjWhCO7x"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Callable\n",
        "from tf_agents.drivers.tf_driver import TFDriver\n",
        "from tf_agents.drivers.dynamic_step_driver import DynamicStepDriver\n",
        "from tf_agents.drivers.dynamic_episode_driver import DynamicEpisodeDriver\n",
        "from tf_agents.environments.tf_py_environment import TFPyEnvironment\n",
        "from tf_agents.metrics.tf_metrics import AverageReturnMetric\n",
        "from tf_agents.policies.tf_policy import TFPolicy\n",
        "from tf_agents.replay_buffers.tf_uniform_replay_buffer import TFUniformReplayBuffer\n",
        "from tf_agents.utils.common import Checkpointer, function\n",
        "from tf_agents.agents.tf_agent import TFAgent"
      ],
      "metadata": {
        "id": "0pScfjPpCwMi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# import tensorflow as tf\n",
        "# from abc import ABC, abstractmethod\n",
        "# from typing import Callable\n",
        "# from tf_agents.drivers.tf_driver import TFDriver\n",
        "# from tf_agents.drivers.dynamic_step_driver import DynamicStepDriver\n",
        "# from tf_agents.drivers.dynamic_episode_driver import DynamicEpisodeDriver\n",
        "# from tf_agents.environments.tf_py_environment import TFPyEnvironment\n",
        "# from tf_agents.metrics.tf_metrics import AverageReturnMetric\n",
        "# from tf_agents.policies.tf_policy import TFPolicy\n",
        "# from tf_agents.replay_buffers.tf_uniform_replay_buffer import TFUniformReplayBuffer\n",
        "# from tf_agents.utils.common import Checkpointer, function\n",
        "# from tf_agents.agents.tf_agent import TFAgent\n",
        "# from agents.agent import Agent\n",
        "\n",
        "\n",
        "class TFAgentBase(Agent, ABC):\n",
        "    def __init__(\n",
        "            self,\n",
        "            agent: TFAgent,\n",
        "            checkpoint_filepath: str,\n",
        "            env_batch_size: int,\n",
        "            replay_memory_capacity: int,\n",
        "            replay_memory_batch_size: int,\n",
        "            trajectory_num_steps: int,\n",
        "            clear_memory_after_train_iteration: bool\n",
        "    ):\n",
        "        self._agent = agent\n",
        "        self._checkpointer = Checkpointer(\n",
        "            ckpt_dir=checkpoint_filepath,\n",
        "            max_to_keep=1,\n",
        "            agent=agent,\n",
        "            policy=agent.policy\n",
        "        )\n",
        "\n",
        "        self._env_batch_size = env_batch_size\n",
        "        self._replay_memory_capacity = replay_memory_capacity\n",
        "        self._replay_memory_batch_size = replay_memory_batch_size\n",
        "        self._trajectory_num_steps = trajectory_num_steps\n",
        "        self._clear_memory_after_train_iteration = clear_memory_after_train_iteration\n",
        "\n",
        "    @property\n",
        "    def train_step_counter(self) -> tf.Variable:\n",
        "        return self._agent.train_step_counter\n",
        "\n",
        "    @property\n",
        "    def policy(self):\n",
        "        return self._agent.policy\n",
        "\n",
        "    @property\n",
        "    def collect_policy(self):\n",
        "        return self._agent.collect_policy\n",
        "\n",
        "    def initialize(self):\n",
        "        self._agent.initialize()\n",
        "        self._checkpointer.initialize_or_restore()\n",
        "\n",
        "    def save(self):\n",
        "        self._checkpointer.save(self._agent.train_step_counter.value())\n",
        "\n",
        "    def load(self):\n",
        "        self._checkpointer.initialize_or_restore()\n",
        "\n",
        "    def _get_replay_memory(self) -> TFUniformReplayBuffer:\n",
        "        return TFUniformReplayBuffer(\n",
        "            data_spec=self._agent.collect_data_spec,\n",
        "            batch_size=self._env_batch_size,\n",
        "            max_length=self._replay_memory_capacity\n",
        "        )\n",
        "\n",
        "    def _get_replay_memory_dataset(self, replay_memory: TFUniformReplayBuffer) -> tf.data.Dataset:\n",
        "        num_steps = self._trajectory_num_steps + 1\n",
        "\n",
        "        return replay_memory.as_dataset(\n",
        "            sample_batch_size=self._replay_memory_batch_size,\n",
        "            num_steps=num_steps,\n",
        "            num_parallel_calls=num_steps + 1,\n",
        "        ).prefetch(num_steps + 1)\n",
        "\n",
        "    def _get_step_driver(\n",
        "            self,\n",
        "            env: TFPyEnvironment,\n",
        "            policy: TFPolicy,\n",
        "            observers: list,\n",
        "            num_steps: int\n",
        "    ):\n",
        "        return DynamicStepDriver(\n",
        "            env=env,\n",
        "            policy=policy,\n",
        "            observers=observers,\n",
        "            num_steps=num_steps\n",
        "        )\n",
        "\n",
        "    def _get_episode_driver(\n",
        "            self,\n",
        "            env: TFPyEnvironment,\n",
        "            policy: TFPolicy,\n",
        "            observers: list,\n",
        "            num_episodes: int\n",
        "    ):\n",
        "        return DynamicEpisodeDriver(env=env, policy=policy, observers=observers, num_episodes=num_episodes)\n",
        "\n",
        "    @abstractmethod\n",
        "    def _get_collect_driver(self, train_env: TFPyEnvironment, observers: list) -> TFDriver:\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def _get_train_step_fn(self) -> Callable:\n",
        "        pass\n",
        "\n",
        "    def train(\n",
        "            self,\n",
        "            train_env: TFPyEnvironment,\n",
        "            eval_env: TFPyEnvironment,\n",
        "            train_iterations: int,\n",
        "            eval_episodes: int,\n",
        "            iterations_per_eval: int,\n",
        "            iterations_per_log: int,\n",
        "            iterations_per_checkpoint: int,\n",
        "            save_best_only: bool\n",
        "    ) -> list:\n",
        "        assert train_env.batch_size == self._env_batch_size, \\\n",
        "            f'AssertionError: Expected environment batch size to be {self._env_batch_size}, got {train_env.batch_size}'\n",
        "\n",
        "        replay_memory = TFUniformReplayBuffer(\n",
        "            data_spec=self._agent.collect_data_spec,\n",
        "            batch_size=self._env_batch_size,\n",
        "            max_length=self._replay_memory_capacity\n",
        "        )\n",
        "        dataset = replay_memory.as_dataset(\n",
        "            sample_batch_size=self._replay_memory_batch_size,\n",
        "            num_steps=None if self._trajectory_num_steps is None else self._trajectory_num_steps + 1,\n",
        "            num_parallel_calls=self._trajectory_num_steps\n",
        "        ).prefetch(1 if self._trajectory_num_steps is None else self._trajectory_num_steps)\n",
        "        collect_driver = self._get_collect_driver(train_env=train_env, observers=[replay_memory.add_batch])\n",
        "        collect_driver.run = function(collect_driver.run)\n",
        "\n",
        "        return self._train(\n",
        "            train_env=train_env,\n",
        "            eval_env=eval_env,\n",
        "            train_iterations=train_iterations,\n",
        "            eval_episodes=eval_episodes,\n",
        "            iterations_per_eval=iterations_per_eval,\n",
        "            iterations_per_log=iterations_per_log,\n",
        "            iterations_per_checkpoint=iterations_per_checkpoint,\n",
        "            replay_memory=replay_memory,\n",
        "            replay_memory_dataset=dataset,\n",
        "            train_collect_driver=collect_driver,\n",
        "            save_best_only=save_best_only\n",
        "        )\n",
        "\n",
        "    def _train(\n",
        "            self,\n",
        "            train_env: TFPyEnvironment,\n",
        "            eval_env: TFPyEnvironment,\n",
        "            train_iterations: int,\n",
        "            eval_episodes: int,\n",
        "            iterations_per_eval: int,\n",
        "            iterations_per_log: int,\n",
        "            iterations_per_checkpoint: int,\n",
        "            replay_memory: TFUniformReplayBuffer,\n",
        "            replay_memory_dataset: tf.data.Dataset,\n",
        "            train_collect_driver: TFDriver,\n",
        "            save_best_only: bool\n",
        "    ) -> list:\n",
        "        average_returns = []\n",
        "        max_avg_reward = -np.inf\n",
        "\n",
        "        eval_metric = AverageReturnMetric(batch_size=eval_env.batch_size, buffer_size=200)\n",
        "        eval_env_driver = self._get_episode_driver(\n",
        "            env=eval_env,\n",
        "            policy=self._agent.policy,\n",
        "            observers=[eval_metric],\n",
        "            num_episodes=eval_episodes\n",
        "        )\n",
        "        dataset_iter = iter(replay_memory_dataset)\n",
        "\n",
        "        eval_env_driver.run = function(eval_env_driver.run)\n",
        "        self._agent.train = function(self._agent.train)\n",
        "        train_step = self._get_train_step_fn()\n",
        "\n",
        "        print('Training has started...')\n",
        "\n",
        "        eval_env_driver.run()\n",
        "        avg_reward = eval_metric.result().numpy()\n",
        "        average_returns.append(avg_reward)\n",
        "        eval_metric.reset()\n",
        "\n",
        "        if max_avg_reward < avg_reward:\n",
        "            max_avg_reward = avg_reward\n",
        "\n",
        "            if save_best_only:\n",
        "                print(\n",
        "                    f'\\nNew best average return found at {max_avg_reward}! '\n",
        "                    f'Saving checkpoint at iteration {0}'\n",
        "                )\n",
        "                self.save()\n",
        "\n",
        "        for i in range(1, train_iterations + 1):\n",
        "            train_collect_driver.run()\n",
        "            train_loss = train_step(replay_memory=replay_memory, dataset_iter=dataset_iter)\n",
        "\n",
        "            if self._clear_memory_after_train_iteration:\n",
        "                replay_memory.clear()\n",
        "\n",
        "            if i % iterations_per_eval == 0:\n",
        "                eval_env_driver.run()\n",
        "                avg_reward = eval_metric.result().numpy()\n",
        "                average_returns.append(avg_reward)\n",
        "                eval_metric.reset()\n",
        "\n",
        "                if max_avg_reward < avg_reward:\n",
        "                    max_avg_reward = avg_reward\n",
        "\n",
        "                    if save_best_only:\n",
        "                        print(\n",
        "                            f'\\nNew best average return found at {max_avg_reward}! '\n",
        "                            f'Saving checkpoint at iteration {i}'\n",
        "                        )\n",
        "                        self.save()\n",
        "\n",
        "            if not save_best_only and i % iterations_per_checkpoint == 0:\n",
        "                print(f'\\nSaving checkpoint at iteration {i}')\n",
        "                self.save()\n",
        "\n",
        "            if i % iterations_per_log == 0:\n",
        "                print(f'\\nIteration: {i}'\n",
        "                      f'\\nTrain Loss: {train_loss}'\n",
        "                      f'\\nAverage Return: {avg_reward}')\n",
        "        return average_returns\n",
        "\n",
        "    def eval(self, eval_env: TFPyEnvironment, num_episodes: int) -> float:\n",
        "        eval_metric = AverageReturnMetric(batch_size=eval_env.batch_size, buffer_size=200)\n",
        "\n",
        "        eval_env_driver = self._get_episode_driver(\n",
        "            env=eval_env,\n",
        "            policy=self._agent.policy,\n",
        "            observers=[eval_metric],\n",
        "            num_episodes=num_episodes\n",
        "        )\n",
        "        eval_env_driver.run = function(eval_env_driver.run)\n",
        "        eval_env_driver.run()\n",
        "        return eval_metric.result().numpy()\n",
        "\n",
        "    def _get_action_step(self, time_step, policy_state=None, use_greedy_policy: bool = True):\n",
        "        policy = self._agent.policy if use_greedy_policy else self._agent.collect_policy\n",
        "\n",
        "        if policy_state is None:\n",
        "            policy_state = policy.get_initial_state(batch_size=self._env_batch_size)\n",
        "        return policy.action(time_step=time_step, policy_state=policy_state)\n",
        "\n",
        "    def compute_action(self, time_step, policy_state=None, use_greedy_policy: bool = True) -> int:\n",
        "        return self._get_action_step(\n",
        "            time_step=time_step,\n",
        "            policy_state=policy_state,\n",
        "            use_greedy_policy=use_greedy_policy\n",
        "        ).action\n",
        "\n",
        "    def get_action_probabilities(self, time_step, policy_state=None, use_greedy_policy: bool = True) -> np.ndarray:\n",
        "        return self._get_action_step(\n",
        "            time_step=time_step,\n",
        "            policy_state=policy_state,\n",
        "            use_greedy_policy=use_greedy_policy\n",
        "        ).info\n"
      ],
      "metadata": {
        "id": "e8wyC4UaCzLf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1#. agents.tfagents.networks - networks.lstm_encoding_network.py\n",
        "\n",
        "---\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "-pUqsRbT5fF0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# coding=utf-8\n",
        "# Copyright 2020 The TF-Agents Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\"\"\"Keras LSTM Encoding Network.\n",
        "\n",
        "Implements a network that will generate the following layers:\n",
        "\n",
        "  [optional]: preprocessing_layers  # preprocessing_layers\n",
        "  [optional]: (Add | Concat(axis=-1) | ...)  # preprocessing_combiner\n",
        "  [optional]: Conv2D # input_conv_layer_params\n",
        "  Flatten\n",
        "  [optional]: Dense  # input_fc_layer_params\n",
        "  [optional]: LSTM cell\n",
        "  [optional]: Dense  # output_fc_layer_params\n",
        "\"\"\"\n",
        "\n",
        "# from __future__ import absolute_import\n",
        "# from __future__ import division\n",
        "# from __future__ import print_function\n",
        "\n",
        "# import gin\n",
        "# import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import\n",
        "\n",
        "# from tf_agents.keras_layers import dynamic_unroll_layer\n",
        "# from tf_agents.networks import encoding_network\n",
        "# from tf_agents.networks import network\n",
        "# from tf_agents.specs import tensor_spec\n",
        "# from tf_agents.trajectories import time_step\n",
        "# from tf_agents.utils import nest_utils\n",
        "\n",
        "KERAS_LSTM_FUSED = 2\n",
        "\n",
        "\n",
        "@gin.configurable\n",
        "class LSTMEncodingNetwork(network.Network):\n",
        "    \"\"\"Recurrent network.\"\"\"\n",
        "\n",
        "    def __init__(\n",
        "            self,\n",
        "            input_tensor_spec,\n",
        "            preprocessing_layers=None,\n",
        "            preprocessing_combiner=None,\n",
        "            conv_layer_params=None,\n",
        "            input_fc_layer_params=(75, 40),\n",
        "            lstm_size=None,\n",
        "            output_fc_layer_params=(75, 40),\n",
        "            activation_fn=tf.keras.activations.relu,\n",
        "            rnn_construction_fn=None,\n",
        "            rnn_construction_kwargs=None,\n",
        "            conv_type='1d',\n",
        "            dtype=tf.float32,\n",
        "            name='LSTMEncodingNetwork',\n",
        "    ):\n",
        "        \"\"\"Creates an instance of `LSTMEncodingNetwork`.\n",
        "\n",
        "        Input preprocessing is possible via `preprocessing_layers` and\n",
        "        `preprocessing_combiner` Layers.  If the `preprocessing_layers` nest is\n",
        "        shallower than `input_tensor_spec`, then the layers will get the subnests.\n",
        "        For example, if:\n",
        "\n",
        "        ```python\n",
        "        input_tensor_spec = ([TensorSpec(3)] * 2, [TensorSpec(3)] * 5)\n",
        "        preprocessing_layers = (Layer1(), Layer2())\n",
        "        ```\n",
        "\n",
        "        then preprocessing will call:\n",
        "\n",
        "        ```python\n",
        "        preprocessed = [preprocessing_layers[0](observations[0]),\n",
        "                        preprocessing_layers[1](observations[1])]\n",
        "        ```\n",
        "\n",
        "        However if\n",
        "\n",
        "        ```python\n",
        "        preprocessing_layers = ([Layer1() for _ in range(2)],\n",
        "                                [Layer2() for _ in range(5)])\n",
        "        ```\n",
        "\n",
        "        then preprocessing will call:\n",
        "        ```python\n",
        "        preprocessed = [\n",
        "          layer(obs) for layer, obs in zip(flatten(preprocessing_layers),\n",
        "                                           flatten(observations))\n",
        "        ]\n",
        "        ```\n",
        "\n",
        "        Args:\n",
        "          input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the\n",
        "            observations.\n",
        "          preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`\n",
        "            representing preprocessing for the different observations. All of these\n",
        "            layers must not be already built.\n",
        "          preprocessing_combiner: (Optional.) A keras layer that takes a flat list\n",
        "            of tensors and combines them.  Good options include\n",
        "            `tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`. This\n",
        "            layer must not be already built.\n",
        "          conv_layer_params: Optional list of convolution layers parameters, where\n",
        "            each item is a length-three tuple indicating (filters, kernel_size,\n",
        "            stride).\n",
        "          input_fc_layer_params: Optional list of fully connected parameters, where\n",
        "            each item is the number of units in the layer. These feed into the\n",
        "            recurrent layer.\n",
        "          lstm_size: An iterable of ints specifying the LSTM cell sizes to use.\n",
        "          output_fc_layer_params: Optional list of fully connected parameters, where\n",
        "            each item is the number of units in the layer. These are applied on top\n",
        "            of the recurrent layer.\n",
        "          activation_fn: Activation function, e.g. tf.keras.activations.relu,.\n",
        "          rnn_construction_fn: (Optional.) Alternate RNN construction function, e.g.\n",
        "            tf.keras.layers.LSTM, tf.keras.layers.CuDNNLSTM. It is invalid to\n",
        "            provide both rnn_construction_fn and lstm_size.\n",
        "          rnn_construction_kwargs: (Optional.) Dictionary or arguments to pass to\n",
        "            rnn_construction_fn.\n",
        "\n",
        "            The RNN will be constructed via:\n",
        "\n",
        "            ```\n",
        "            rnn_layer = rnn_construction_fn(**rnn_construction_kwargs)\n",
        "            ```\n",
        "          dtype: The dtype to use by the convolution, LSTM, and fully connected\n",
        "            layers.\n",
        "          name: A string representing name of the network.\n",
        "\n",
        "        Raises:\n",
        "          ValueError: If any of `preprocessing_layers` is already built.\n",
        "          ValueError: If `preprocessing_combiner` is already built.\n",
        "          ValueError: If neither `lstm_size` nor `rnn_construction_fn` are provided.\n",
        "          ValueError: If both `lstm_size` and `rnn_construction_fn` are provided.\n",
        "        \"\"\"\n",
        "        if lstm_size is None and rnn_construction_fn is None:\n",
        "            raise ValueError('Need to provide either custom rnn_construction_fn or '\n",
        "                             'lstm_size.')\n",
        "        if lstm_size and rnn_construction_fn:\n",
        "            raise ValueError('Cannot provide both custom rnn_construction_fn and '\n",
        "                             'lstm_size.')\n",
        "\n",
        "        kernel_initializer = tf.compat.v1.variance_scaling_initializer(\n",
        "            scale=2.0, mode='fan_in', distribution='truncated_normal')\n",
        "\n",
        "        input_encoder = encoding_network.EncodingNetwork(\n",
        "            input_tensor_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            conv_layer_params=conv_layer_params,\n",
        "            fc_layer_params=input_fc_layer_params,\n",
        "            activation_fn=activation_fn,\n",
        "            kernel_initializer=kernel_initializer,\n",
        "            conv_type=conv_type,\n",
        "            dtype=dtype)\n",
        "\n",
        "        # Create RNN cell\n",
        "        if rnn_construction_fn:\n",
        "            rnn_construction_kwargs = rnn_construction_kwargs or {}\n",
        "            lstm_network = rnn_construction_fn(**rnn_construction_kwargs)\n",
        "        else:\n",
        "            if len(lstm_size) == 1:\n",
        "                cell = tf.keras.layers.LSTMCell(\n",
        "                    lstm_size[0],\n",
        "                    dtype=dtype,\n",
        "                    implementation=KERAS_LSTM_FUSED)\n",
        "            else:\n",
        "                cell = tf.keras.layers.StackedRNNCells(\n",
        "                    [tf.keras.layers.LSTMCell(size, dtype=dtype,\n",
        "                                              implementation=KERAS_LSTM_FUSED)\n",
        "                     for size in lstm_size])\n",
        "            lstm_network = dynamic_unroll_layer.DynamicUnroll(cell)\n",
        "\n",
        "        output_encoder = []\n",
        "        if output_fc_layer_params:\n",
        "            output_encoder = [\n",
        "                tf.keras.layers.Dense(\n",
        "                    num_units,\n",
        "                    activation=activation_fn,\n",
        "                    kernel_initializer=kernel_initializer,\n",
        "                    dtype=dtype) for num_units in output_fc_layer_params\n",
        "            ]\n",
        "\n",
        "        counter = [-1]\n",
        "\n",
        "        def create_spec(size):\n",
        "            counter[0] += 1\n",
        "            return tensor_spec.TensorSpec(\n",
        "                size, dtype=dtype, name='network_state_%d' % counter[0])\n",
        "\n",
        "        state_spec = tf.nest.map_structure(create_spec,\n",
        "                                           lstm_network.cell.state_size)\n",
        "\n",
        "        super(LSTMEncodingNetwork, self).__init__(\n",
        "            input_tensor_spec=input_tensor_spec, state_spec=state_spec, name=name)\n",
        "\n",
        "        self._conv_layer_params = conv_layer_params\n",
        "        self._input_encoder = input_encoder\n",
        "        self._lstm_network = lstm_network\n",
        "        self._output_encoder = output_encoder\n",
        "\n",
        "    def call(self,\n",
        "             observation,\n",
        "             step_type,\n",
        "             network_state=(),\n",
        "             training=False):\n",
        "        \"\"\"Apply the network.\n",
        "\n",
        "        Args:\n",
        "          observation: A tuple of tensors matching `input_tensor_spec`.\n",
        "          step_type: A tensor of `StepType.\n",
        "          network_state: (optional.) The network state.\n",
        "          training: Whether the output is being used for training.\n",
        "\n",
        "        Returns:\n",
        "          `(outputs, network_state)` - the network output and next network state.\n",
        "\n",
        "        Raises:\n",
        "          ValueError: If observation tensors lack outer `(batch,)` or\n",
        "            `(batch, time)` axes.\n",
        "        \"\"\"\n",
        "        num_outer_dims = nest_utils.get_outer_rank(observation,\n",
        "                                                   self.input_tensor_spec)\n",
        "        if num_outer_dims not in (1, 2):\n",
        "            raise ValueError(\n",
        "                'Input observation must have a batch or batch x time outer shape.')\n",
        "\n",
        "        has_time_dim = num_outer_dims == 2\n",
        "        if not has_time_dim:\n",
        "            # Add a time dimension to the inputs.\n",
        "            observation = tf.nest.map_structure(lambda t: tf.expand_dims(t, 1),\n",
        "                                                observation)\n",
        "            step_type = tf.nest.map_structure(lambda t: tf.expand_dims(t, 1),\n",
        "                                              step_type)\n",
        "\n",
        "        state, _ = self._input_encoder(\n",
        "            observation, step_type=step_type, network_state=(), training=training)\n",
        "\n",
        "        network_kwargs = {}\n",
        "        if isinstance(self._lstm_network, dynamic_unroll_layer.DynamicUnroll):\n",
        "            network_kwargs['reset_mask'] = tf.equal(step_type,\n",
        "                                                    time_step.StepType.FIRST,\n",
        "                                                    name='mask')\n",
        "\n",
        "        # Unroll over the time sequence.\n",
        "        output = self._lstm_network(\n",
        "            inputs=state,\n",
        "            initial_state=network_state,\n",
        "            training=training,\n",
        "            **network_kwargs)\n",
        "\n",
        "        if isinstance(self._lstm_network, dynamic_unroll_layer.DynamicUnroll):\n",
        "            state, network_state = output\n",
        "        else:\n",
        "            state = output[0]\n",
        "            network_state = tf.nest.pack_sequence_as(\n",
        "                self._lstm_network.cell.state_size, tf.nest.flatten(output[1:]))\n",
        "\n",
        "        for layer in self._output_encoder:\n",
        "            state = layer(state, training=training)\n",
        "\n",
        "        if not has_time_dim:\n",
        "            # Remove time dimension from the state.\n",
        "            state = tf.squeeze(state, [1])\n",
        "\n",
        "        return state, network_state\n"
      ],
      "metadata": {
        "id": "GpZBHQUoZZu5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1A. agents.tfagents.networks - networks.actor_distribution_network.py\n"
      ],
      "metadata": {
        "id": "UmYDaVud1CaG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# coding=utf-8\n",
        "# Copyright 2020 The TF-Agents Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\"\"\"Sample Keras actor network that generates distributions.\"\"\"\n",
        "\n",
        "# from __future__ import absolute_import\n",
        "# from __future__ import division\n",
        "# from __future__ import print_function\n",
        "\n",
        "# import gin\n",
        "# import numpy as np\n",
        "# import tensorflow as tf\n",
        "# import tensorflow_probability as tfp\n",
        "\n",
        "# from tf_agents.networks import categorical_projection_network\n",
        "# from tf_agents.networks import encoding_network\n",
        "# from tf_agents.networks import network\n",
        "# from tf_agents.networks import normal_projection_network\n",
        "# from tf_agents.specs import tensor_spec\n",
        "# from tf_agents.utils import nest_utils\n",
        "\n",
        "\n",
        "def _categorical_projection_net(action_spec, logits_init_output_factor=0.1):\n",
        "    return categorical_projection_network.CategoricalProjectionNetwork(\n",
        "        action_spec, logits_init_output_factor=logits_init_output_factor)\n",
        "\n",
        "\n",
        "def _normal_projection_net(action_spec,\n",
        "                           init_action_stddev=0.35,\n",
        "                           init_means_output_factor=0.1,\n",
        "                           seed_stream_class=tfp.util.SeedStream,\n",
        "                           seed=None):\n",
        "    std_bias_initializer_value = np.log(np.exp(init_action_stddev) - 1)\n",
        "\n",
        "    return normal_projection_network.NormalProjectionNetwork(\n",
        "        action_spec,\n",
        "        init_means_output_factor=init_means_output_factor,\n",
        "        std_bias_initializer_value=std_bias_initializer_value,\n",
        "        scale_distribution=False,\n",
        "        seed_stream_class=seed_stream_class,\n",
        "        seed=seed)\n",
        "\n",
        "\n",
        "@gin.configurable\n",
        "class ActorDistributionNetwork(network.DistributionNetwork):\n",
        "    \"\"\"Creates an actor producing either Normal or Categorical distribution.\n",
        "    Note: By default, this network uses `NormalProjectionNetwork` for continuous\n",
        "    projection which by default uses `tanh_squash_to_spec` to normalize its\n",
        "    output. Due to the nature of the `tanh` function, values near the spec bounds\n",
        "    cannot be returned.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 input_tensor_spec,\n",
        "                 output_tensor_spec,\n",
        "                 preprocessing_layers=None,\n",
        "                 preprocessing_combiner=None,\n",
        "                 conv_layer_params=None,\n",
        "                 fc_layer_params=(200, 100),\n",
        "                 dropout_layer_params=None,\n",
        "                 activation_fn=tf.keras.activations.relu,\n",
        "                 kernel_initializer=None,\n",
        "                 seed_stream_class=tfp.util.SeedStream,\n",
        "                 seed=None,\n",
        "                 batch_squash=True,\n",
        "                 dtype=tf.float32,\n",
        "                 discrete_projection_net=_categorical_projection_net,\n",
        "                 continuous_projection_net=_normal_projection_net,\n",
        "                 conv_type='1d',\n",
        "                 name='ActorDistributionNetwork'):\n",
        "        \"\"\"Creates an instance of `ActorDistributionNetwork`.\n",
        "        Args:\n",
        "          input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the\n",
        "            input.\n",
        "          output_tensor_spec: A nest of `tensor_spec.BoundedTensorSpec` representing\n",
        "            the output.\n",
        "          preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`\n",
        "            representing preprocessing for the different observations.\n",
        "            All of these layers must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          preprocessing_combiner: (Optional.) A keras layer that takes a flat list\n",
        "            of tensors and combines them. Good options include\n",
        "            `tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.\n",
        "            This layer must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          conv_layer_params: Optional list of convolution layers parameters, where\n",
        "            each item is a length-three tuple indicating (filters, kernel_size,\n",
        "            stride).\n",
        "          fc_layer_params: Optional list of fully_connected parameters, where each\n",
        "            item is the number of units in the layer.\n",
        "          dropout_layer_params: Optional list of dropout layer parameters, each item\n",
        "            is the fraction of input units to drop or a dictionary of parameters\n",
        "            according to the keras.Dropout documentation. The additional parameter\n",
        "            `permanent`, if set to True, allows to apply dropout at inference for\n",
        "            approximated Bayesian inference. The dropout layers are interleaved with\n",
        "            the fully connected layers; there is a dropout layer after each fully\n",
        "            connected layer, except if the entry in the list is None. This list must\n",
        "            have the same length of fc_layer_params, or be None.\n",
        "          activation_fn: Activation function, e.g. tf.nn.relu, slim.leaky_relu, ...\n",
        "          kernel_initializer: Initializer to use for the kernels of the conv and\n",
        "            dense layers. If none is provided a default glorot_uniform.\n",
        "          seed_stream_class: The seed stream class. This is almost always\n",
        "            tfp.util.SeedStream, except for in unit testing, when one may want to\n",
        "            seed all the layers deterministically.\n",
        "          seed: seed used for Keras kernal initializers for NormalProjectionNetwork.\n",
        "          batch_squash: If True the outer_ranks of the observation are squashed into\n",
        "            the batch dimension. This allow encoding networks to be used with\n",
        "            observations with shape [BxTx...].\n",
        "          dtype: The dtype to use by the convolution and fully connected layers.\n",
        "          discrete_projection_net: Callable that generates a discrete projection\n",
        "            network to be called with some hidden state and the outer_rank of the\n",
        "            state.\n",
        "          continuous_projection_net: Callable that generates a continuous projection\n",
        "            network to be called with some hidden state and the outer_rank of the\n",
        "            state.\n",
        "          name: A string representing name of the network.\n",
        "        Raises:\n",
        "          ValueError: If `input_tensor_spec` contains more than one observation.\n",
        "        \"\"\"\n",
        "\n",
        "        if not kernel_initializer:\n",
        "            kernel_initializer = tf.compat.v1.keras.initializers.glorot_uniform()\n",
        "\n",
        "        encoder = encoding_network.EncodingNetwork(\n",
        "            input_tensor_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            conv_layer_params=conv_layer_params,\n",
        "            fc_layer_params=fc_layer_params,\n",
        "            dropout_layer_params=dropout_layer_params,\n",
        "            activation_fn=activation_fn,\n",
        "            kernel_initializer=kernel_initializer,\n",
        "            batch_squash=batch_squash,\n",
        "            conv_type=conv_type,\n",
        "            dtype=dtype)\n",
        "\n",
        "        def map_proj(spec):\n",
        "            if tensor_spec.is_discrete(spec):\n",
        "                return discrete_projection_net(spec)\n",
        "            else:\n",
        "                kwargs = {}\n",
        "                if continuous_projection_net is _normal_projection_net:\n",
        "                    kwargs['seed'] = seed\n",
        "                    kwargs['seed_stream_class'] = seed_stream_class\n",
        "                return continuous_projection_net(spec, **kwargs)\n",
        "\n",
        "        projection_networks = tf.nest.map_structure(map_proj, output_tensor_spec)\n",
        "        output_spec = tf.nest.map_structure(lambda proj_net: proj_net.output_spec,\n",
        "                                            projection_networks)\n",
        "\n",
        "        super(ActorDistributionNetwork, self).__init__(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            state_spec=(),\n",
        "            output_spec=output_spec,\n",
        "            name=name)\n",
        "\n",
        "        self._encoder = encoder\n",
        "        self._projection_networks = projection_networks\n",
        "        self._output_tensor_spec = output_tensor_spec\n",
        "\n",
        "    @property\n",
        "    def output_tensor_spec(self):\n",
        "        return self._output_tensor_spec\n",
        "\n",
        "    def call(self,\n",
        "             observations,\n",
        "             step_type,\n",
        "             network_state,\n",
        "             training=False,\n",
        "             mask=None):\n",
        "        state, network_state = self._encoder(\n",
        "            observations,\n",
        "            step_type=step_type,\n",
        "            network_state=network_state,\n",
        "            training=training)\n",
        "        outer_rank = nest_utils.get_outer_rank(observations, self.input_tensor_spec)\n",
        "\n",
        "        def call_projection_net(proj_net):\n",
        "            distribution, _ = proj_net(\n",
        "                state, outer_rank, training=training, mask=mask)\n",
        "            return distribution\n",
        "\n",
        "        output_actions = tf.nest.map_structure(\n",
        "            call_projection_net, self._projection_networks)\n",
        "        return output_actions, network_state"
      ],
      "metadata": {
        "id": "8HHtsFrYZZsQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1B. agents.tfagents.networks - networks.actor_distribution_rnn_network.py\n"
      ],
      "metadata": {
        "id": "Uee33HN67ZUA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# coding=utf-8\n",
        "# Copyright 2020 The TF-Agents Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\"\"\"Sample Keras actor network  with LSTM cells that generates distributions.\"\"\"\n",
        "\n",
        "# from __future__ import absolute_import\n",
        "# from __future__ import division\n",
        "# from __future__ import print_function\n",
        "\n",
        "\n",
        "# import gin\n",
        "# import numpy as np\n",
        "# import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import\n",
        "\n",
        "# from tf_agents.networks import categorical_projection_network\n",
        "# from tf_agents.networks import network\n",
        "# from tf_agents.networks import normal_projection_network\n",
        "# from tf_agents.specs import tensor_spec\n",
        "# from tf_agents.utils import nest_utils\n",
        "\n",
        "# # from agents.tfagents.networks.lstm_encoding_network import LSTMEncodingNetwork   THIS WAS CALLED IN THE NOTEBOOK JUST EARLIER\n",
        "\n",
        "\n",
        "def _categorical_projection_net(action_spec, logits_init_output_factor=0.1):\n",
        "    return categorical_projection_network.CategoricalProjectionNetwork(\n",
        "        action_spec, logits_init_output_factor=logits_init_output_factor)\n",
        "\n",
        "\n",
        "def _normal_projection_net(action_spec,\n",
        "                           init_action_stddev=0.35,\n",
        "                           init_means_output_factor=0.1):\n",
        "    std_bias_initializer_value = np.log(np.exp(init_action_stddev) - 1)\n",
        "\n",
        "    return normal_projection_network.NormalProjectionNetwork(\n",
        "        action_spec,\n",
        "        init_means_output_factor=init_means_output_factor,\n",
        "        std_bias_initializer_value=std_bias_initializer_value)\n",
        "\n",
        "\n",
        "@gin.configurable\n",
        "class ActorDistributionRnnNetwork(network.DistributionNetwork):\n",
        "    \"\"\"Creates an actor producing either Normal or Categorical distribution.\n",
        "    Note: By default, this network uses `NormalProjectionNetwork` for continuous\n",
        "    projection which by default uses `tanh_squash_to_spec` to normalize its\n",
        "    output. Due to the nature of the `tanh` function, values near the spec bounds\n",
        "    cannot be returned.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 input_tensor_spec,\n",
        "                 output_tensor_spec,\n",
        "                 preprocessing_layers=None,\n",
        "                 preprocessing_combiner=None,\n",
        "                 conv_layer_params=None,\n",
        "                 input_fc_layer_params=(200, 100),\n",
        "                 input_dropout_layer_params=None,\n",
        "                 lstm_size=None,\n",
        "                 output_fc_layer_params=(200, 100),\n",
        "                 activation_fn=tf.keras.activations.relu,\n",
        "                 dtype=tf.float32,\n",
        "                 discrete_projection_net=_categorical_projection_net,\n",
        "                 continuous_projection_net=_normal_projection_net,\n",
        "                 rnn_construction_fn=None,\n",
        "                 rnn_construction_kwargs={},\n",
        "                 conv_type='1d',\n",
        "                 name='ActorDistributionRnnNetwork'):\n",
        "        \"\"\"Creates an instance of `ActorDistributionRnnNetwork`.\n",
        "        Args:\n",
        "          input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the\n",
        "            input.\n",
        "          output_tensor_spec: A nest of `tensor_spec.BoundedTensorSpec` representing\n",
        "            the output.\n",
        "          preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`\n",
        "            representing preprocessing for the different observations.\n",
        "            All of these layers must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          preprocessing_combiner: (Optional.) A keras layer that takes a flat list\n",
        "            of tensors and combines them. Good options include\n",
        "            `tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.\n",
        "            This layer must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          conv_layer_params: Optional list of convolution layers parameters, where\n",
        "            each item is a length-three tuple indicating (filters, kernel_size,\n",
        "            stride).\n",
        "          input_fc_layer_params: Optional list of fully_connected parameters, where\n",
        "            each item is the number of units in the layer. This is applied before\n",
        "            the LSTM cell.\n",
        "          input_dropout_layer_params: Optional list of dropout layer parameters,\n",
        "            each item is the fraction of input units to drop or a dictionary of\n",
        "            parameters according to the keras.Dropout documentation. The additional\n",
        "            parameter `permanent`, if set to True, allows to apply dropout at\n",
        "            inference for approximated Bayesian inference. The dropout layers are\n",
        "            interleaved with the fully connected layers; there is a dropout layer\n",
        "            after each fully connected layer, except if the entry in the list is\n",
        "            None. This list must have the same length of input_fc_layer_params, or\n",
        "            be None.\n",
        "          lstm_size: An iterable of ints specifying the LSTM cell sizes to use.\n",
        "          output_fc_layer_params: Optional list of fully_connected parameters, where\n",
        "            each item is the number of units in the layer. This is applied after the\n",
        "            LSTM cell.\n",
        "          activation_fn: Activation function, e.g. tf.nn.relu, slim.leaky_relu, ...\n",
        "          dtype: The dtype to use by the convolution and fully connected layers.\n",
        "          discrete_projection_net: Callable that generates a discrete projection\n",
        "            network to be called with some hidden state and the outer_rank of the\n",
        "            state.\n",
        "          continuous_projection_net: Callable that generates a continuous projection\n",
        "            network to be called with some hidden state and the outer_rank of the\n",
        "            state.\n",
        "          rnn_construction_fn: (Optional.) Alternate RNN construction function, e.g.\n",
        "            tf.keras.layers.LSTM, tf.keras.layers.CuDNNLSTM. It is invalid to\n",
        "            provide both rnn_construction_fn and lstm_size.\n",
        "          rnn_construction_kwargs: (Optional.) Dictionary or arguments to pass to\n",
        "            rnn_construction_fn.\n",
        "            The RNN will be constructed via:\n",
        "            ```\n",
        "            rnn_layer = rnn_construction_fn(**rnn_construction_kwargs)\n",
        "            ```\n",
        "          name: A string representing name of the network.\n",
        "        Raises:\n",
        "          ValueError: If 'input_dropout_layer_params' is not None.\n",
        "        \"\"\"\n",
        "        if input_dropout_layer_params:\n",
        "            raise ValueError('Dropout layer is not supported.')\n",
        "\n",
        "        lstm_encoder = LSTMEncodingNetwork(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            conv_layer_params=conv_layer_params,\n",
        "            input_fc_layer_params=input_fc_layer_params,\n",
        "            lstm_size=lstm_size,\n",
        "            output_fc_layer_params=output_fc_layer_params,\n",
        "            activation_fn=activation_fn,\n",
        "            rnn_construction_fn=rnn_construction_fn,\n",
        "            rnn_construction_kwargs=rnn_construction_kwargs,\n",
        "            conv_type=conv_type,\n",
        "            dtype=dtype,\n",
        "            name=name)\n",
        "\n",
        "        def map_proj(spec):\n",
        "            if tensor_spec.is_discrete(spec):\n",
        "                return discrete_projection_net(spec)\n",
        "            else:\n",
        "                return continuous_projection_net(spec)\n",
        "\n",
        "        projection_networks = tf.nest.map_structure(map_proj, output_tensor_spec)\n",
        "        output_spec = tf.nest.map_structure(lambda proj_net: proj_net.output_spec,\n",
        "                                            projection_networks)\n",
        "\n",
        "        super(ActorDistributionRnnNetwork, self).__init__(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            state_spec=lstm_encoder.state_spec,\n",
        "            output_spec=output_spec,\n",
        "            name=name)\n",
        "\n",
        "        self._lstm_encoder = lstm_encoder\n",
        "        self._projection_networks = projection_networks\n",
        "        self._output_tensor_spec = output_tensor_spec\n",
        "\n",
        "    @property\n",
        "    def output_tensor_spec(self):\n",
        "        return self._output_tensor_spec\n",
        "\n",
        "    def call(self, observation, step_type, network_state=(), training=False):\n",
        "        state, network_state = self._lstm_encoder(\n",
        "            observation, step_type=step_type, network_state=network_state,\n",
        "            training=training)\n",
        "        outer_rank = nest_utils.get_outer_rank(observation, self.input_tensor_spec)\n",
        "        output_actions = tf.nest.map_structure(\n",
        "            lambda proj_net: proj_net(state, outer_rank, training=training)[0],\n",
        "            self._projection_networks)\n",
        "        return output_actions, network_state"
      ],
      "metadata": {
        "id": "QiSjmwt8ZZxc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1C. agents.tfagents.networks - networks.value_network.py"
      ],
      "metadata": {
        "id": "KTRBwHLz76ve"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# coding=utf-8\n",
        "# Copyright 2020 The TF-Agents Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\"\"\"Sample Keras Value Network.\n",
        "Implements a network that will generate the following layers:\n",
        "  [optional]: preprocessing_layers  # preprocessing_layers\n",
        "  [optional]: (Add | Concat(axis=-1) | ...)  # preprocessing_combiner\n",
        "  [optional]: Conv2D # conv_layer_params\n",
        "  Flatten\n",
        "  [optional]: Dense  # fc_layer_params\n",
        "  Dense -> 1         # Value output\n",
        "\"\"\"\n",
        "\n",
        "# from __future__ import absolute_import\n",
        "# from __future__ import division\n",
        "# from __future__ import print_function\n",
        "\n",
        "# import gin\n",
        "# import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import\n",
        "\n",
        "# from tf_agents.networks import encoding_network\n",
        "# from tf_agents.networks import network\n",
        "\n",
        "\n",
        "@gin.configurable\n",
        "class ValueNetwork(network.Network):\n",
        "    \"\"\"Feed Forward value network. Reduces to 1 value output per batch item.\"\"\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 input_tensor_spec,\n",
        "                 preprocessing_layers=None,\n",
        "                 preprocessing_combiner=None,\n",
        "                 conv_layer_params=None,\n",
        "                 fc_layer_params=(75, 40),\n",
        "                 dropout_layer_params=None,\n",
        "                 activation_fn=tf.keras.activations.relu,\n",
        "                 kernel_initializer=None,\n",
        "                 batch_squash=True,\n",
        "                 dtype=tf.float32,\n",
        "                 conv_type='1d',\n",
        "                 name='ValueNetwork'):\n",
        "        \"\"\"Creates an instance of `ValueNetwork`.\n",
        "        Network supports calls with shape outer_rank + observation_spec.shape. Note\n",
        "        outer_rank must be at least 1.\n",
        "        Args:\n",
        "          input_tensor_spec: A `tensor_spec.TensorSpec` or a tuple of specs\n",
        "            representing the input observations.\n",
        "          preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`\n",
        "            representing preprocessing for the different observations.\n",
        "            All of these layers must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          preprocessing_combiner: (Optional.) A keras layer that takes a flat list\n",
        "            of tensors and combines them. Good options include\n",
        "            `tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.\n",
        "            This layer must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          conv_layer_params: Optional list of convolution layers parameters, where\n",
        "            each item is a length-three tuple indicating (filters, kernel_size,\n",
        "            stride).\n",
        "          fc_layer_params: Optional list of fully_connected parameters, where each\n",
        "            item is the number of units in the layer.\n",
        "          dropout_layer_params: Optional list of dropout layer parameters, each item\n",
        "            is the fraction of input units to drop or a dictionary of parameters\n",
        "            according to the keras.Dropout documentation. The additional parameter\n",
        "            `permanent`, if set to True, allows to apply dropout at inference for\n",
        "            approximated Bayesian inference. The dropout layers are interleaved with\n",
        "            the fully connected layers; there is a dropout layer after each fully\n",
        "            connected layer, except if the entry in the list is None. This list must\n",
        "            have the same length of fc_layer_params, or be None.\n",
        "          activation_fn: Activation function, e.g. tf.keras.activations.relu,.\n",
        "          kernel_initializer: Initializer to use for the kernels of the conv and\n",
        "            dense layers. If none is provided a default variance_scaling_initializer\n",
        "          batch_squash: If True the outer_ranks of the observation are squashed into\n",
        "            the batch dimension. This allow encoding networks to be used with\n",
        "            observations with shape [BxTx...].\n",
        "          dtype: The dtype to use by the convolution and fully connected layers.\n",
        "          name: A string representing name of the network.\n",
        "        Raises:\n",
        "          ValueError: If input_tensor_spec is not an instance of network.InputSpec.\n",
        "        \"\"\"\n",
        "        super(ValueNetwork, self).__init__(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            state_spec=(),\n",
        "            name=name)\n",
        "\n",
        "        if not kernel_initializer:\n",
        "            kernel_initializer = tf.compat.v1.keras.initializers.glorot_uniform()\n",
        "\n",
        "        self._encoder = encoding_network.EncodingNetwork(\n",
        "            input_tensor_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            conv_layer_params=conv_layer_params,\n",
        "            fc_layer_params=fc_layer_params,\n",
        "            dropout_layer_params=dropout_layer_params,\n",
        "            activation_fn=activation_fn,\n",
        "            kernel_initializer=kernel_initializer,\n",
        "            batch_squash=batch_squash,\n",
        "            conv_type=conv_type,\n",
        "            dtype=dtype)\n",
        "\n",
        "        self._postprocessing_layers = tf.keras.layers.Dense(\n",
        "            1,\n",
        "            activation=None,\n",
        "            kernel_initializer=tf.random_uniform_initializer(\n",
        "                minval=-0.03, maxval=0.03))\n",
        "\n",
        "    def call(self, observation, step_type=None, network_state=(), training=False):\n",
        "        state, network_state = self._encoder(\n",
        "            observation, step_type=step_type, network_state=network_state,\n",
        "            training=training)\n",
        "        value = self._postprocessing_layers(state, training=training)\n",
        "        return tf.squeeze(value, -1), network_state"
      ],
      "metadata": {
        "id": "HK71FybB-zom"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1D. agents.tfagents.networks - networks.value_rnn_network.py\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "VhZNJpCg-sFe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# coding=utf-8\n",
        "# Copyright 2020 The TF-Agents Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\"\"\"Sample Keras Value Network with LSTM cells .\n",
        "Implements a network that will generate the following layers:\n",
        "  [optional]: preprocessing_layers  # preprocessing_layers\n",
        "  [optional]: (Add | Concat(axis=-1) | ...)  # preprocessing_combiner\n",
        "  [optional]: Conv2D # conv_layer_params\n",
        "  Flatten\n",
        "  [optional]: Dense  # input_fc_layer_params\n",
        "  [optional]: LSTM   # lstm_cell_params\n",
        "  [optional]: Dense  # output_fc_layer_params\n",
        "  Dense -> 1         # Value output\n",
        "\"\"\"\n",
        "\n",
        "# from __future__ import absolute_import\n",
        "# from __future__ import division\n",
        "# from __future__ import print_function\n",
        "\n",
        "# import gin\n",
        "# import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import\n",
        "\n",
        "# from tf_agents.networks import network\n",
        "\n",
        "# # from agents.tfagents.networks.lstm_encoding_network import LSTMEncodingNetwork\n",
        "\n",
        "\n",
        "@gin.configurable\n",
        "class ValueRnnNetwork(network.Network):\n",
        "    \"\"\"Recurrent value network. Reduces to 1 value output per batch item.\"\"\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 input_tensor_spec,\n",
        "                 preprocessing_layers=None,\n",
        "                 preprocessing_combiner=None,\n",
        "                 conv_layer_params=None,\n",
        "                 input_fc_layer_params=(75, 40),\n",
        "                 input_dropout_layer_params=None,\n",
        "                 lstm_size=(40,),\n",
        "                 output_fc_layer_params=(75, 40),\n",
        "                 activation_fn=tf.keras.activations.relu,\n",
        "                 dtype=tf.float32,\n",
        "                 conv_type='1d',\n",
        "                 name='ValueRnnNetwork'):\n",
        "        \"\"\"Creates an instance of `ValueRnnNetwork`.\n",
        "        Network supports calls with shape outer_rank + input_tensor_shape.shape.\n",
        "        Note outer_rank must be at least 1.\n",
        "        Args:\n",
        "          input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the\n",
        "            input observations.\n",
        "          preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`\n",
        "            representing preprocessing for the different observations.\n",
        "            All of these layers must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          preprocessing_combiner: (Optional.) A keras layer that takes a flat list\n",
        "            of tensors and combines them.  Good options include\n",
        "            `tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.\n",
        "            This layer must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          conv_layer_params: Optional list of convolution layers parameters, where\n",
        "            each item is a length-three tuple indicating (filters, kernel_size,\n",
        "            stride).\n",
        "          input_fc_layer_params: Optional list of fully_connected parameters, where\n",
        "            each item is the number of units in the layer. This is applied before\n",
        "            the LSTM cell.\n",
        "          input_dropout_layer_params: Optional list of dropout layer parameters,\n",
        "            where each item is the fraction of input units to drop. The dropout\n",
        "            layers are interleaved with the fully connected layers; there is a\n",
        "            dropout layer after each fully connected layer, except if the entry in\n",
        "            the list is None. This list must have the same length of\n",
        "            input_fc_layer_params, or be None.\n",
        "          lstm_size: An iterable of ints specifying the LSTM cell sizes to use.\n",
        "          output_fc_layer_params: Optional list of fully_connected parameters, where\n",
        "            each item is the number of units in the layer. This is applied after the\n",
        "            LSTM cell.\n",
        "          activation_fn: Activation function, e.g. tf.keras.activations.relu,.\n",
        "          dtype: The dtype to use by the convolution, LSTM, and fully connected\n",
        "            layers.\n",
        "          name: A string representing name of the network.\n",
        "        \"\"\"\n",
        "        del input_dropout_layer_params\n",
        "\n",
        "        lstm_encoder = LSTMEncodingNetwork(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            conv_layer_params=conv_layer_params,\n",
        "            input_fc_layer_params=input_fc_layer_params,\n",
        "            lstm_size=lstm_size,\n",
        "            output_fc_layer_params=output_fc_layer_params,\n",
        "            activation_fn=activation_fn,\n",
        "            dtype=dtype,\n",
        "            conv_type=conv_type,\n",
        "            name=name)\n",
        "\n",
        "        postprocessing_layers = tf.keras.layers.Dense(\n",
        "            1,\n",
        "            activation=None,\n",
        "            kernel_initializer=tf.random_uniform_initializer(\n",
        "                minval=-0.03, maxval=0.03))\n",
        "\n",
        "        super(ValueRnnNetwork, self).__init__(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            state_spec=lstm_encoder.state_spec,\n",
        "            name=name)\n",
        "\n",
        "        self._lstm_encoder = lstm_encoder\n",
        "        self._postprocessing_layers = postprocessing_layers\n",
        "\n",
        "    def call(self,\n",
        "             observation,\n",
        "             step_type=None,\n",
        "             network_state=(),\n",
        "             training=False):\n",
        "        state, network_state = self._lstm_encoder(\n",
        "            observation, step_type=step_type, network_state=network_state,\n",
        "            training=training)\n",
        "        value = self._postprocessing_layers(state, training=training)\n",
        "        return tf.squeeze(value, -1), network_state"
      ],
      "metadata": {
        "id": "rToAEJHU76cE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 1. agents.tfagents - tfagents.ppo.py\n"
      ],
      "metadata": {
        "id": "3w0OybrJDjlk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "from typing import Callable\n",
        "from tf_agents.agents.ppo.ppo_clip_agent import PPOClipAgent\n",
        "from tf_agents.drivers.tf_driver import TFDriver\n",
        "from tf_agents.environments.tf_py_environment import TFPyEnvironment\n",
        "from tf_agents.replay_buffers.tf_uniform_replay_buffer import TFUniformReplayBuffer\n",
        "from tf_agents.utils.common import function"
      ],
      "metadata": {
        "id": "AuGBddELDj08"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import math\n",
        "# import tensorflow as tf\n",
        "# from typing import Callable\n",
        "# from tf_agents.agents.ppo.ppo_clip_agent import PPOClipAgent\n",
        "# from tf_agents.drivers.tf_driver import TFDriver\n",
        "# from tf_agents.environments.tf_py_environment import TFPyEnvironment\n",
        "# from tf_agents.replay_buffers.tf_uniform_replay_buffer import TFUniformReplayBuffer\n",
        "# from tf_agents.utils.common import function\n",
        "# # from agents.tfagents.tfagent import TFAgentBase\n",
        "# # from agents.tfagents.networks.actor_distribution_network import ActorDistributionNetwork\n",
        "# # from agents.tfagents.networks.actor_distribution_rnn_network import ActorDistributionRnnNetwork\n",
        "# # from agents.tfagents.networks.value_network import ValueNetwork\n",
        "# # from agents.tfagents.networks.value_rnn_network import ValueRnnNetwork\n",
        "\n",
        "\n",
        "class PPOAgent(TFAgentBase):\n",
        "    def __init__(\n",
        "            self,\n",
        "            input_tensor_spec,\n",
        "            action_spec,\n",
        "            time_step_spec,\n",
        "            env_batch_size: int,\n",
        "            checkpoint_filepath: str,\n",
        "            fc_layers: list or None,\n",
        "            conv_layers: list[tuple[int, int, int]] or None = None,\n",
        "            conv_type: str = '1d',\n",
        "            lstm_layers: list[int] or None = None,\n",
        "            train_sequence_len: int = 1,\n",
        "            preprocessing_layers: list or dict = None,\n",
        "            preprocessing_combiner: list or dict or Callable = None,\n",
        "            greedy_eval: bool = True,\n",
        "            epsilon_clipping: float = 0.3,\n",
        "            lambda_value: float = 0.95,\n",
        "            entropy_regularization_init: float = 0.0,\n",
        "            entropy_regularization_min: float = 0.0,\n",
        "            entropy_regularization_decay_steps: int = 0,\n",
        "            num_epochs: int = 40,\n",
        "            use_gae: bool = True,\n",
        "            gamma: float = 0.99,\n",
        "            replay_memory_capacity: int = 1000,\n",
        "            replay_memory_batch_size: int or None = None,\n",
        "            collection_episodes_per_iteration: int = 5,\n",
        "            train_step_counter: tf.Variable = tf.Variable(initial_value=0)\n",
        "    ):\n",
        "        assert train_sequence_len >= 1, \\\n",
        "            f'AssertionError: train_sequence_len is expected to be >= 1, got {train_sequence_len}'\n",
        "        assert lstm_layers is None or train_sequence_len > 1, \\\n",
        "            'AssertionError: train_sequence_len is expected to be greater than 1 if lstm_layers is not None'\n",
        "\n",
        "        assert entropy_regularization_decay_steps == 0 or entropy_regularization_init == 0.0 or \\\n",
        "               entropy_regularization_min < entropy_regularization_init <= 1.0, \\\n",
        "            'AssertionError: entropy_regularization_decay_steps is expected to be zero ' \\\n",
        "            'or entropy_regularization_init is expected to be zero or ' \\\n",
        "            'or entropy_regularization_min < entropy_regularization_init <= 1.0' \\\n",
        "            f'got entropy_regularization_decay_steps={entropy_regularization_decay_steps}, ' \\\n",
        "            f'entropy_regularization_init={entropy_regularization_init}, ' \\\n",
        "            f'entropy_regularization_min={entropy_regularization_min}'\n",
        "\n",
        "        if entropy_regularization_decay_steps > 0:\n",
        "            entropy_regularization_fn = tf.keras.optimizers.schedules.PolynomialDecay(\n",
        "                initial_learning_rate=entropy_regularization_init,\n",
        "                decay_steps=entropy_regularization_decay_steps,\n",
        "                end_learning_rate=entropy_regularization_min\n",
        "            )\n",
        "            entropy_regularization = lambda: entropy_regularization_fn(train_step_counter)\n",
        "        else:\n",
        "            entropy_regularization = entropy_regularization_init\n",
        "\n",
        "        if lstm_layers is None:\n",
        "            actor_network = ActorDistributionNetwork(\n",
        "                input_tensor_spec=input_tensor_spec,\n",
        "                output_tensor_spec=action_spec,\n",
        "                preprocessing_layers=preprocessing_layers,\n",
        "                preprocessing_combiner=preprocessing_combiner,\n",
        "                fc_layer_params=fc_layers,\n",
        "                conv_layer_params=conv_layers,\n",
        "                conv_type=conv_type,\n",
        "                activation_fn='gelu'\n",
        "            )\n",
        "            value_network = ValueNetwork(\n",
        "                input_tensor_spec=input_tensor_spec,\n",
        "                preprocessing_layers=preprocessing_layers,\n",
        "                preprocessing_combiner=preprocessing_combiner,\n",
        "                fc_layer_params=fc_layers,\n",
        "                conv_layer_params=conv_layers,\n",
        "                conv_type=conv_type,\n",
        "                activation_fn='gelu'\n",
        "            )\n",
        "        else:\n",
        "            actor_network = ActorDistributionRnnNetwork(\n",
        "                input_tensor_spec=input_tensor_spec,\n",
        "                output_tensor_spec=action_spec,\n",
        "                preprocessing_layers=preprocessing_layers,\n",
        "                preprocessing_combiner=preprocessing_combiner,\n",
        "                lstm_size=lstm_layers,\n",
        "                output_fc_layer_params=fc_layers,\n",
        "                conv_layer_params=conv_layers,\n",
        "                conv_type=conv_type,\n",
        "                activation_fn='gelu'\n",
        "            )\n",
        "            value_network = ValueRnnNetwork(\n",
        "                input_tensor_spec=input_tensor_spec,\n",
        "                preprocessing_layers=preprocessing_layers,\n",
        "                preprocessing_combiner=preprocessing_combiner,\n",
        "                lstm_size=lstm_layers,\n",
        "                output_fc_layer_params=fc_layers,\n",
        "                conv_layer_params=conv_layers,\n",
        "                conv_type=conv_type,\n",
        "                activation_fn='gelu'\n",
        "            )\n",
        "\n",
        "        agent = PPOClipAgent(\n",
        "            time_step_spec=time_step_spec,\n",
        "            action_spec=action_spec,\n",
        "            optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),\n",
        "            actor_net=actor_network,\n",
        "            value_net=value_network,\n",
        "            greedy_eval=greedy_eval,\n",
        "            importance_ratio_clipping=epsilon_clipping,\n",
        "            lambda_value=lambda_value,\n",
        "            discount_factor=gamma,\n",
        "            entropy_regularization=entropy_regularization,\n",
        "            num_epochs=num_epochs,\n",
        "            use_gae=use_gae,\n",
        "            train_step_counter=train_step_counter\n",
        "        )\n",
        "\n",
        "        super().__init__(\n",
        "            agent=agent,\n",
        "            checkpoint_filepath=checkpoint_filepath,\n",
        "            env_batch_size=env_batch_size,\n",
        "            replay_memory_capacity=replay_memory_capacity,\n",
        "            replay_memory_batch_size=replay_memory_batch_size,\n",
        "            trajectory_num_steps=train_sequence_len,\n",
        "            clear_memory_after_train_iteration=True\n",
        "        )\n",
        "\n",
        "        self._replay_memory_batch_size = replay_memory_batch_size\n",
        "        self._collection_episodes_per_iteration = collection_episodes_per_iteration\n",
        "\n",
        "        self._optimized_minibatch_train_fn = function(self._train_step_minibatch_fn)\n",
        "\n",
        "    def _get_collect_driver(self, train_env: TFPyEnvironment, observers: list) -> TFDriver:\n",
        "        return super()._get_episode_driver(\n",
        "            env=train_env,\n",
        "            policy=self.collect_policy,\n",
        "            observers=observers,\n",
        "            num_episodes=self._collection_episodes_per_iteration\n",
        "        )\n",
        "\n",
        "    def _train_step_batch_fn(self, replay_memory: TFUniformReplayBuffer, dataset_iter: iter) -> float:\n",
        "        trajectories = replay_memory.gather_all()\n",
        "        return self._agent.train(experience=trajectories).loss\n",
        "\n",
        "    def _train_step_minibatch_fn(self, dataset_iter: iter) -> float:\n",
        "        trajectories, _ = next(dataset_iter)\n",
        "        return self._agent.train(experience=trajectories).loss\n",
        "\n",
        "    def _train_step_minibatch(self, replay_memory: TFUniformReplayBuffer, dataset_iter: iter) -> float:\n",
        "        train_loss = 0\n",
        "        optimized_minibatch_train_fn = self._optimized_minibatch_train_fn\n",
        "\n",
        "        num_memory_items = replay_memory.num_frames()\n",
        "        num_batches = math.ceil(num_memory_items / self._replay_memory_batch_size)\n",
        "\n",
        "        for _ in range(num_batches):\n",
        "            train_loss += optimized_minibatch_train_fn(dataset_iter=dataset_iter)\n",
        "        return train_loss / num_batches\n",
        "\n",
        "    def _get_train_step_fn(self) -> Callable:\n",
        "        if self._replay_memory_batch_size is None:\n",
        "            return function(self._train_step_batch_fn)\n",
        "        else:\n",
        "            return self._train_step_minibatch\n"
      ],
      "metadata": {
        "id": "mLBJ2G1BDycn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2A. agents.tfagents.networks - networks.q_network.py\n"
      ],
      "metadata": {
        "id": "1YeeylmY_UXn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# coding=utf-8\n",
        "# Copyright 2020 The TF-Agents Authors.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "\"\"\"Sample Keras networks for DQN.\"\"\"\n",
        "\n",
        "# from __future__ import absolute_import\n",
        "# from __future__ import division\n",
        "# from __future__ import print_function\n",
        "\n",
        "# import gin\n",
        "# import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import\n",
        "\n",
        "# from tf_agents.networks import encoding_network\n",
        "# from tf_agents.networks import network\n",
        "\n",
        "\n",
        "def validate_specs(action_spec, observation_spec):\n",
        "    \"\"\"Validates the spec contains a single action.\"\"\"\n",
        "    del observation_spec  # not currently validated\n",
        "\n",
        "    flat_action_spec = tf.nest.flatten(action_spec)\n",
        "    if len(flat_action_spec) > 1:\n",
        "        raise ValueError('Network only supports action_specs with a single action.')\n",
        "\n",
        "    if flat_action_spec[0].shape not in [(), (1,)]:\n",
        "        raise ValueError(\n",
        "            'Network only supports action_specs with shape in [(), (1,)])')\n",
        "\n",
        "\n",
        "@gin.configurable\n",
        "class QNetwork(network.Network):\n",
        "    \"\"\"Feed Forward network.\"\"\"\n",
        "\n",
        "    def __init__(self,\n",
        "                 input_tensor_spec,\n",
        "                 action_spec,\n",
        "                 preprocessing_layers=None,\n",
        "                 preprocessing_combiner=None,\n",
        "                 conv_layer_params=None,\n",
        "                 fc_layer_params=(75, 40),\n",
        "                 dropout_layer_params=None,\n",
        "                 activation_fn=tf.keras.activations.relu,\n",
        "                 kernel_initializer=None,\n",
        "                 batch_squash=True,\n",
        "                 dtype=tf.float32,\n",
        "                 conv_type='1d',\n",
        "                 q_layer_activation_fn=None,\n",
        "                 name='QNetwork'):\n",
        "        \"\"\"Creates an instance of `QNetwork`.\n",
        "        Args:\n",
        "          input_tensor_spec: A nest of `tensor_spec.TensorSpec` representing the\n",
        "            input observations.\n",
        "          action_spec: A nest of `tensor_spec.BoundedTensorSpec` representing the\n",
        "            actions.\n",
        "          preprocessing_layers: (Optional.) A nest of `tf.keras.layers.Layer`\n",
        "            representing preprocessing for the different observations.\n",
        "            All of these layers must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          preprocessing_combiner: (Optional.) A keras layer that takes a flat list\n",
        "            of tensors and combines them. Good options include\n",
        "            `tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.\n",
        "            This layer must not be already built. For more details see\n",
        "            the documentation of `networks.EncodingNetwork`.\n",
        "          conv_layer_params: Optional list of convolution layers parameters, where\n",
        "            each item is a length-three tuple indicating (filters, kernel_size,\n",
        "            stride).\n",
        "          fc_layer_params: Optional list of fully_connected parameters, where each\n",
        "            item is the number of units in the layer.\n",
        "          dropout_layer_params: Optional list of dropout layer parameters, where\n",
        "            each item is the fraction of input units to drop. The dropout layers are\n",
        "            interleaved with the fully connected layers; there is a dropout layer\n",
        "            after each fully connected layer, except if the entry in the list is\n",
        "            None. This list must have the same length of fc_layer_params, or be\n",
        "            None.\n",
        "          activation_fn: Activation function, e.g. tf.keras.activations.relu.\n",
        "          kernel_initializer: Initializer to use for the kernels of the conv and\n",
        "            dense layers. If none is provided a default variance_scaling_initializer\n",
        "          batch_squash: If True the outer_ranks of the observation are squashed into\n",
        "            the batch dimension. This allow encoding networks to be used with\n",
        "            observations with shape [BxTx...].\n",
        "          dtype: The dtype to use by the convolution and fully connected layers.\n",
        "          q_layer_activation_fn: Activation function for the Q layer.\n",
        "          name: A string representing the name of the network.\n",
        "        Raises:\n",
        "          ValueError: If `input_tensor_spec` contains more than one observation. Or\n",
        "            if `action_spec` contains more than one action.\n",
        "        \"\"\"\n",
        "        validate_specs(action_spec, input_tensor_spec)\n",
        "        action_spec = tf.nest.flatten(action_spec)[0]\n",
        "        num_actions = action_spec.maximum - action_spec.minimum + 1\n",
        "        encoder_input_tensor_spec = input_tensor_spec\n",
        "\n",
        "        encoder = encoding_network.EncodingNetwork(\n",
        "            encoder_input_tensor_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            conv_layer_params=conv_layer_params,\n",
        "            fc_layer_params=fc_layer_params,\n",
        "            dropout_layer_params=dropout_layer_params,\n",
        "            activation_fn=activation_fn,\n",
        "            kernel_initializer=kernel_initializer,\n",
        "            batch_squash=batch_squash,\n",
        "            conv_type=conv_type,\n",
        "            dtype=dtype)\n",
        "\n",
        "        q_value_layer = tf.keras.layers.Dense(\n",
        "            num_actions,\n",
        "            activation=q_layer_activation_fn,\n",
        "            kernel_initializer=tf.random_uniform_initializer(\n",
        "                minval=-0.03, maxval=0.03),\n",
        "            bias_initializer=tf.constant_initializer(-0.2),\n",
        "            dtype=dtype)\n",
        "\n",
        "        super(QNetwork, self).__init__(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            state_spec=(),\n",
        "            name=name)\n",
        "\n",
        "        self._encoder = encoder\n",
        "        self._q_value_layer = q_value_layer\n",
        "\n",
        "    def call(self, observation, step_type=None, network_state=(), training=False):\n",
        "        \"\"\"Runs the given observation through the network.\n",
        "        Args:\n",
        "          observation: The observation to provide to the network.\n",
        "          step_type: The step type for the given observation. See `StepType` in\n",
        "            time_step.py.\n",
        "          network_state: A state tuple to pass to the network, mainly used by RNNs.\n",
        "          training: Whether the output is being used for training.\n",
        "        Returns:\n",
        "          A tuple `(logits, network_state)`.\n",
        "        \"\"\"\n",
        "        state, network_state = self._encoder(\n",
        "            observation, step_type=step_type, network_state=network_state,\n",
        "            training=training)\n",
        "        q_value = self._q_value_layer(state, training=training)\n",
        "        return q_value, network_state"
      ],
      "metadata": {
        "id": "6tUkgQ4j_TeY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2. agents.tfagents - tfagents.dqn.py"
      ],
      "metadata": {
        "id": "CP7ztxOWBgIb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Callable\n",
        "from tf_agents.agents.dqn import dqn_agent\n",
        "from tf_agents.drivers.tf_driver import TFDriver\n",
        "from tf_agents.environments.tf_py_environment import TFPyEnvironment\n",
        "from tf_agents.policies.random_tf_policy import RandomTFPolicy\n",
        "from tf_agents.replay_buffers.tf_uniform_replay_buffer import TFUniformReplayBuffer\n",
        "from tf_agents.utils.common import function"
      ],
      "metadata": {
        "id": "HpZuh2pOEhcL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import tensorflow as tf\n",
        "# from typing import Callable\n",
        "# from tf_agents.agents.dqn import dqn_agent\n",
        "# from tf_agents.drivers.tf_driver import TFDriver\n",
        "# from tf_agents.environments.tf_py_environment import TFPyEnvironment\n",
        "# from tf_agents.policies.random_tf_policy import RandomTFPolicy\n",
        "# from tf_agents.replay_buffers.tf_uniform_replay_buffer import TFUniformReplayBuffer\n",
        "# from tf_agents.utils.common import function\n",
        "# # from agents.tfagents.tfagent import TFAgentBase\n",
        "# # from agents.tfagents.networks.q_network import QNetwork\n",
        "\n",
        "\n",
        "class DQNAgent(TFAgentBase):\n",
        "    def __init__(\n",
        "            self,\n",
        "            input_tensor_spec,\n",
        "            action_spec,\n",
        "            time_step_spec,\n",
        "            env_batch_size: int,\n",
        "            checkpoint_filepath: str,\n",
        "            fc_layers: list or None,\n",
        "            conv_layers: list[tuple[int, int, int]] or None = None,\n",
        "            conv_type: str = '1d',\n",
        "            preprocessing_layers: list or dict = None,\n",
        "            preprocessing_combiner: list or dict or Callable = None,\n",
        "            n_step: int = 3,\n",
        "            double_dqn: bool = True,\n",
        "            tau: float = 0.005,\n",
        "            target_update_steps: int = 1,\n",
        "            epsilon_init: float = 0.1,\n",
        "            epsilon_min: float = 0.1,\n",
        "            epsilon_decay_steps: int = 0,\n",
        "            gamma: float = 0.99,\n",
        "            replay_memory_capacity: int = 50000,\n",
        "            replay_memory_batch_size: int = 64,\n",
        "            initial_collect_steps: int = 10000,\n",
        "            collection_steps_per_iteration: int = 1,\n",
        "            train_step_counter: tf.Variable = tf.Variable(initial_value=0)\n",
        "    ):\n",
        "        assert epsilon_decay_steps == 0 or \\\n",
        "               epsilon_decay_steps > 0 and epsilon_init is not None and epsilon_min < epsilon_init <= 1.0, \\\n",
        "            'AssertionError: epsilon_decay_steps is expected to be zero or epsilon_min < epsilon_init <= 1.0' \\\n",
        "            f'got epsilon_decay_steps={epsilon_decay_steps}, epsilon_init={epsilon_init}, epsilon_min={epsilon_min}'\n",
        "\n",
        "        if epsilon_decay_steps > 0:\n",
        "            epsilon_greedy_fn = tf.keras.optimizers.schedules.PolynomialDecay(\n",
        "                initial_learning_rate=epsilon_init,\n",
        "                decay_steps=epsilon_decay_steps,\n",
        "                end_learning_rate=epsilon_min\n",
        "            )\n",
        "            epsilon_greedy = lambda: epsilon_greedy_fn(train_step_counter)\n",
        "        else:\n",
        "            epsilon_greedy = epsilon_init\n",
        "\n",
        "        q_network = QNetwork(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            action_spec=action_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            fc_layer_params=fc_layers,\n",
        "            conv_layer_params=conv_layers,\n",
        "            conv_type=conv_type,\n",
        "            activation_fn='gelu'\n",
        "        )\n",
        "        target_network = QNetwork(\n",
        "            input_tensor_spec=input_tensor_spec,\n",
        "            action_spec=action_spec,\n",
        "            preprocessing_layers=preprocessing_layers,\n",
        "            preprocessing_combiner=preprocessing_combiner,\n",
        "            fc_layer_params=fc_layers,\n",
        "            conv_layer_params=conv_layers,\n",
        "            conv_type=conv_type,\n",
        "            activation_fn='gelu'\n",
        "        )\n",
        "\n",
        "        if double_dqn:\n",
        "            agent = dqn_agent.DdqnAgent(\n",
        "                time_step_spec=time_step_spec,\n",
        "                action_spec=action_spec,\n",
        "                q_network=q_network,\n",
        "                target_q_network=target_network,\n",
        "                optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),\n",
        "                n_step_update=n_step,\n",
        "                target_update_tau=tau,\n",
        "                target_update_period=target_update_steps,\n",
        "                gamma=gamma,\n",
        "                epsilon_greedy=epsilon_greedy,\n",
        "                train_step_counter=train_step_counter\n",
        "            )\n",
        "        else:\n",
        "            agent = dqn_agent.DqnAgent(\n",
        "                time_step_spec=time_step_spec,\n",
        "                action_spec=action_spec,\n",
        "                q_network=q_network,\n",
        "                target_q_network=target_network,\n",
        "                optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),\n",
        "                n_step_update=n_step,\n",
        "                target_update_tau=tau,\n",
        "                target_update_period=target_update_steps,\n",
        "                gamma=gamma,\n",
        "                epsilon_greedy=epsilon_greedy,\n",
        "                train_step_counter=train_step_counter\n",
        "            )\n",
        "\n",
        "        super().__init__(\n",
        "            agent=agent,\n",
        "            checkpoint_filepath=checkpoint_filepath,\n",
        "            env_batch_size=env_batch_size,\n",
        "            replay_memory_capacity=replay_memory_capacity,\n",
        "            replay_memory_batch_size=replay_memory_batch_size,\n",
        "            trajectory_num_steps=n_step,\n",
        "            clear_memory_after_train_iteration=False\n",
        "        )\n",
        "\n",
        "        self._n_step = n_step\n",
        "        self._initial_collect_steps = initial_collect_steps\n",
        "        self._collection_steps_per_iteration = collection_steps_per_iteration\n",
        "\n",
        "    def _get_collect_driver(self, train_env: TFPyEnvironment, observers: list) -> TFDriver:\n",
        "        return super()._get_step_driver(\n",
        "            env=train_env,\n",
        "            policy=self.collect_policy,\n",
        "            observers=observers,\n",
        "            num_steps=self._collection_steps_per_iteration * self._n_step\n",
        "        )\n",
        "\n",
        "    def _train_step_fn(self, replay_memory: TFUniformReplayBuffer, dataset_iter: iter) -> float:\n",
        "        trajectories, _ = next(dataset_iter)\n",
        "        return self._agent.train(trajectories).loss\n",
        "\n",
        "    def _get_train_step_fn(self) -> Callable:\n",
        "        return function(self._train_step_fn)\n",
        "\n",
        "    def _train(\n",
        "            self,\n",
        "            train_env: TFPyEnvironment,\n",
        "            eval_env: TFPyEnvironment,\n",
        "            train_iterations: int,\n",
        "            eval_episodes: int,\n",
        "            iterations_per_eval: int,\n",
        "            iterations_per_log: int,\n",
        "            iterations_per_checkpoint: int,\n",
        "            replay_memory: TFUniformReplayBuffer,\n",
        "            replay_memory_dataset: tf.data.Dataset,\n",
        "            train_collect_driver: TFDriver,\n",
        "            save_best_only: bool\n",
        "    ) -> list:\n",
        "        print('Collecting Initial Samples...')\n",
        "\n",
        "        initial_collect_policy = RandomTFPolicy(train_env.time_step_spec(), train_env.action_spec())\n",
        "        initial_collect_driver = super()._get_step_driver(\n",
        "            env=train_env,\n",
        "            policy=initial_collect_policy,\n",
        "            observers=[replay_memory.add_batch],\n",
        "            num_steps=self._initial_collect_steps\n",
        "        )\n",
        "        initial_collect_driver.run = function(initial_collect_driver.run)\n",
        "        initial_collect_driver.run()\n",
        "\n",
        "        return super()._train(\n",
        "            train_env=train_env,\n",
        "            eval_env=eval_env,\n",
        "            train_iterations=train_iterations,\n",
        "            eval_episodes=eval_episodes,\n",
        "            iterations_per_eval=iterations_per_eval,\n",
        "            iterations_per_log=iterations_per_log,\n",
        "            iterations_per_checkpoint=iterations_per_checkpoint,\n",
        "            replay_memory=replay_memory,\n",
        "            replay_memory_dataset=replay_memory_dataset,\n",
        "            train_collect_driver=train_collect_driver,\n",
        "            save_best_only=save_best_only\n",
        "        )\n"
      ],
      "metadata": {
        "id": "CfUxfrinEi8d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# environments.actions - environments.actions.py"
      ],
      "metadata": {
        "id": "h0lrsvTcf5ln"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from enum import Enum"
      ],
      "metadata": {
        "id": "dZ4PaxY2gtRU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# from enum import Enum\n",
        "\n",
        "\n",
        "class Action(Enum):\n",
        "    BUY = 0\n",
        "    SELL = 1\n",
        "    HOLD = 2"
      ],
      "metadata": {
        "id": "kGlJua0Gf5B-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# metrics.metric - metrics.metric.py"
      ],
      "metadata": {
        "id": "Fvdgx052geOQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from abc import ABC, abstractmethod\n",
        "\n",
        "\n",
        "class Metric(ABC):\n",
        "    def __init__(self, name: str):\n",
        "        self._name = name\n",
        "        self._episode_metrics = []\n",
        "\n",
        "    @property\n",
        "    def name(self) -> str:\n",
        "        return self._name\n",
        "\n",
        "    @property\n",
        "    def episode_metrics(self) -> list[float]:\n",
        "        return self._episode_metrics\n",
        "\n",
        "    @abstractmethod\n",
        "    def reset(self):\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def update(self, log_pnl: float):\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def result(self) -> float:\n",
        "        pass\n",
        "\n",
        "    def register(self):\n",
        "        self._episode_metrics.append(self.result())\n"
      ],
      "metadata": {
        "id": "c-IkatyJgdnc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# environments.environment - environments.environment.py\n"
      ],
      "metadata": {
        "id": "Jg5aQDQLfrwd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import gym"
      ],
      "metadata": {
        "id": "_K3iKowthI_H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import gym\n",
        "# import numpy as np\n",
        "# from environments.actions import Action\n",
        "# from metrics.metric import Metric\n",
        "\n",
        "class TradingEnvironment(gym.Env):\n",
        "    def __init__(self, env_config: dict):\n",
        "        assert 'states' in env_config, 'AssertionError: Expected \"states\" in env_config'\n",
        "        assert 'reward_fn' in env_config, 'AssertionError: Expected \"reward_function\" in env_config'\n",
        "        assert 'episode_steps' in env_config, 'AssertionError: Expected \"episode_steps\" in env_config'\n",
        "        assert 'metrics' in env_config, 'AssertionError: Expected \"metrics\" in env_config'\n",
        "\n",
        "        self._states = env_config['states']\n",
        "        self._reward_function = env_config['reward_fn']\n",
        "        self._episode_steps = env_config['episode_steps']\n",
        "\n",
        "        self._metrics = env_config['metrics']\n",
        "\n",
        "        if self._metrics is None:\n",
        "            self._metrics = []\n",
        "\n",
        "        self._num_states = self._states.shape[0] - 1\n",
        "\n",
        "        assert self._num_states >= self._episode_steps, \\\n",
        "            'AssertionError: Not enough states are provided in the environment: ' \\\n",
        "            f'num_states = {self._num_states}, episode_steps = {self._episode_steps}'\n",
        "\n",
        "        self._state_index = 0\n",
        "\n",
        "        self.observation_space = gym.spaces.Box(\n",
        "            low=self._states.min(axis=0),\n",
        "            high=self._states.max(axis=0),\n",
        "            shape=self._states.shape[1:],\n",
        "            dtype=np.float32\n",
        "        )\n",
        "        self.action_space = gym.spaces.Discrete(n=len(Action))\n",
        "\n",
        "        assert self._states.dtype == self.observation_space.dtype, \\\n",
        "            f'AssertionError: Expected states to have dtype = {self.observation_space.dtype}, got {self._states.dtype}'\n",
        "\n",
        "    @property\n",
        "    def metrics(self) -> list[Metric]:\n",
        "        return self._metrics\n",
        "\n",
        "    def update_metrics(self, log_pnl: float):\n",
        "        for metric in self._metrics:\n",
        "            metric.update(log_pnl=log_pnl)\n",
        "\n",
        "    def register_metrics(self):\n",
        "        for metric in self._metrics:\n",
        "            metric.register()\n",
        "\n",
        "    def reset(self) -> np.ndarray:\n",
        "        for metric in self._metrics:\n",
        "            metric.reset()\n",
        "        return self._states[self._state_index]\n",
        "\n",
        "    def step(self, action: int) -> (np.ndarray, float, bool):\n",
        "        reward = self._reward_function.get_reward(i=self._state_index, action=action)\n",
        "\n",
        "        self._state_index += 1\n",
        "        next_state = self._states[self._state_index]\n",
        "\n",
        "        if self._state_index == self._num_states:\n",
        "            done = True\n",
        "            self._state_index = 0\n",
        "        elif self._state_index % self._episode_steps == 0:\n",
        "            done = True\n",
        "        else:\n",
        "            done = False\n",
        "\n",
        "        log_pnl = 0.0 if action == Action.HOLD.value else reward\n",
        "        self.update_metrics(log_pnl=log_pnl)\n",
        "\n",
        "        if done:\n",
        "            self.register_metrics()\n",
        "        return next_state, reward, done\n",
        "\n",
        "    def render(self, mode=None):\n",
        "        print('\\n--- Current State ---')\n",
        "        print(self._states[self._state_index])\n"
      ],
      "metadata": {
        "id": "Mo1pgbfYfrJG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# rules.rule import Rule\n"
      ],
      "metadata": {
        "id": "q1g1V5ovcKXU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from abc import ABC, abstractmethod\n",
        "\n",
        "\n",
        "class Rule(ABC):\n",
        "    @abstractmethod\n",
        "    def filter(self, action: int) -> int:\n",
        "        pass"
      ],
      "metadata": {
        "id": "im3sg7hzcKiN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# environments.wrappers.tf.tfenv import TFTradingEnvironment\n"
      ],
      "metadata": {
        "id": "JzkdHmOhh0eT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from tf_agents.environments.py_environment import PyEnvironment\n",
        "from tf_agents.specs.array_spec import BoundedArraySpec\n",
        "from tf_agents.trajectories import time_step"
      ],
      "metadata": {
        "id": "8piQGXvabAM9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from tf_agents.environments.py_environment import PyEnvironment\n",
        "# from tf_agents.specs.array_spec import BoundedArraySpec\n",
        "# from tf_agents.trajectories import time_step\n",
        "# from environments.actions import Action\n",
        "# from environments.environment import TradingEnvironment\n",
        "# from metrics.metric import Metric\n",
        "# from rules.rule import Rule\n",
        "\n",
        "\n",
        "class TFTradingEnvironment(PyEnvironment):\n",
        "    def __init__(self, env: TradingEnvironment):\n",
        "        super().__init__()\n",
        "\n",
        "        self._env = env\n",
        "\n",
        "        self._action_spec = BoundedArraySpec(\n",
        "            shape=(), dtype=np.int32, minimum=0, maximum=len(Action) - 1, name='action'\n",
        "        )\n",
        "        self._observation_spec = BoundedArraySpec(\n",
        "            shape=self._env.observation_space.shape, dtype=self._env.observation_space.dtype,\n",
        "            minimum=self._env.observation_space.low, maximum=self._env.observation_space.high, name='observation'\n",
        "        )\n",
        "        self._done = False\n",
        "        self._discount_rate = 1.0\n",
        "\n",
        "    def get_metrics(self) -> list[Metric]:\n",
        "        return self._env.metrics\n",
        "\n",
        "    def get_episode_metrics(self) -> dict[str, float]:\n",
        "        metrics = self._env.metrics\n",
        "        return {metric.name: metric.result() for metric in metrics}\n",
        "\n",
        "    def action_spec(self) -> BoundedArraySpec:\n",
        "        return self._action_spec\n",
        "\n",
        "    def observation_spec(self) -> BoundedArraySpec:\n",
        "        return self._observation_spec\n",
        "\n",
        "    def _reset(self) -> time_step.TimeStep:\n",
        "        self._done = False\n",
        "        observation = self._env.reset()\n",
        "        return time_step.restart(observation=observation)\n",
        "\n",
        "    def _step(self, action_spec) -> time_step.TimeStep:\n",
        "        if self._done:\n",
        "            return self._reset()\n",
        "\n",
        "        action = action_spec.item()\n",
        "        next_observation, reward, self._done = self._env.step(action)\n",
        "\n",
        "        if self._done:\n",
        "            return time_step.termination(observation=next_observation, reward=reward)\n",
        "        else:\n",
        "            return time_step.transition(observation=next_observation, reward=reward, discount=self._discount_rate)\n",
        "\n",
        "    def render(self, **kwargs):\n",
        "        self._env.render()\n",
        "\n",
        "\n",
        "class TFRuleTradingEnvironment(TFTradingEnvironment):\n",
        "    def __init__(self, env: TradingEnvironment, rules: list[Rule] or None):\n",
        "        super().__init__(env=env)\n",
        "\n",
        "        self._rules = rules\n",
        "\n",
        "    def _step(self, action_spec) -> time_step.TimeStep:\n",
        "        if self._done:\n",
        "            return self._reset()\n",
        "\n",
        "        action = action_spec.item()\n",
        "\n",
        "        if self._rules is not None:\n",
        "            for rule in self._rules:\n",
        "                action = rule.filter(action=action)\n",
        "\n",
        "        next_observation, reward, self._done = self._env.step(action)\n",
        "\n",
        "        if self._done:\n",
        "            return time_step.termination(observation=next_observation, reward=reward)\n",
        "        else:\n",
        "            return time_step.transition(observation=next_observation, reward=reward, discount=self._discount_rate)\n"
      ],
      "metadata": {
        "id": "ad-PIBDGaARW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# environments.rewards.function import RewardFunction"
      ],
      "metadata": {
        "id": "yi9PSbEXbU3U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from abc import ABC, abstractmethod\n",
        "\n",
        "\n",
        "class RewardFunction(ABC):\n",
        "    def __init__(\n",
        "            self,\n",
        "            timeframe_size: int,\n",
        "            target_horizon_len: int,\n",
        "            highs: np.ndarray,\n",
        "            lows: np.ndarray,\n",
        "            closes: np.ndarray,\n",
        "            fees_percentage: float,\n",
        "            verbose: bool = False\n",
        "    ):\n",
        "        rewards_fn = self._build_reward_fn(\n",
        "            timeframe_size=timeframe_size,\n",
        "            target_horizon_len=target_horizon_len,\n",
        "            highs=highs,\n",
        "            lows=lows,\n",
        "            closes=closes\n",
        "        )\n",
        "\n",
        "        fees = np.log((1 - fees_percentage)/(1 + fees_percentage))\n",
        "        rewards_fn[:, 0:2] += fees\n",
        "        hold_rewards = -rewards_fn.max(axis=1)\n",
        "        hold_rewards[hold_rewards > 0] = 0\n",
        "        self._rewards_fn = np.hstack((rewards_fn, np.expand_dims(hold_rewards, axis=-1)))\n",
        "\n",
        "        if verbose:\n",
        "            print(f'Rewards: {self._rewards_fn.shape}')\n",
        "\n",
        "    @property\n",
        "    def reward_fn(self) -> np.ndarray:\n",
        "        return self._rewards_fn\n",
        "\n",
        "    @reward_fn.setter\n",
        "    def reward_fn(self, reward_fn: np.ndarray):\n",
        "        self._rewards_fn = reward_fn\n",
        "\n",
        "    def __call__(self, i: int, action: int) -> float:\n",
        "        return self.get_reward(i=i, action=action)\n",
        "\n",
        "    def get_reward(self, i: int, action: int) -> float:\n",
        "        return self._rewards_fn[i, action]\n",
        "\n",
        "    def get_reward_fn_shape(self):\n",
        "        return self._rewards_fn.shape\n",
        "\n",
        "    @abstractmethod\n",
        "    def _build_reward_fn(\n",
        "            self,\n",
        "            timeframe_size: int,\n",
        "            target_horizon_len: int,\n",
        "            highs: np.ndarray,\n",
        "            lows: np.ndarray,\n",
        "            closes: np.ndarray\n",
        "    ) -> np.ndarray:\n",
        "        pass\n"
      ],
      "metadata": {
        "id": "Imijrj_zbVHf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#  environments.rewards.smurf import SmurfRewardFunction"
      ],
      "metadata": {
        "id": "-SLUmKH6h0lT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from environments.rewards.function import RewardFunction\n",
        "\n",
        "\n",
        "class SmurfRewardFunction:\n",
        "    def __init__(\n",
        "            self,\n",
        "            reward_function: RewardFunction\n",
        "    ):\n",
        "        self._reward_function = reward_function\n",
        "\n",
        "        smurf_rf = self._reward_function.reward_fn\n",
        "        smurf_rf[:, 2] = 0.0055\n",
        "        self._reward_function.reward_fn = smurf_rf\n",
        "        print(self._reward_function.reward_fn[0])\n",
        "\n",
        "    def __call__(self, i: int, action: int) -> float:\n",
        "        return self.get_reward(i=i, action=action)\n",
        "\n",
        "    def get_reward(self, i: int, action: int) -> float:\n",
        "        return self._reward_function.get_reward(i=i, action=action)\n",
        "\n",
        "    def get_reward_fn_shape(self):\n",
        "        return self._reward_function.get_reward_fn_shape()\n"
      ],
      "metadata": {
        "id": "Nbe9rEbJaTPz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# environments.rewards.marketorder import MarketOrderRF\n"
      ],
      "metadata": {
        "id": "rg5djtsNh0uC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from environments.rewards.function import RewardFunction\n",
        "\n",
        "\n",
        "class MarketOrderRF(RewardFunction):\n",
        "    def __init__(\n",
        "            self,\n",
        "            timeframe_size: int,\n",
        "            target_horizon_len: int,\n",
        "            highs: np.ndarray,\n",
        "            lows: np.ndarray,\n",
        "            closes: np.ndarray,\n",
        "            fees_percentage: float\n",
        "    ):\n",
        "        super().__init__(\n",
        "            timeframe_size=timeframe_size,\n",
        "            target_horizon_len=target_horizon_len,\n",
        "            highs=highs,\n",
        "            lows=lows,\n",
        "            closes=closes,\n",
        "            fees_percentage=fees_percentage\n",
        "        )\n",
        "\n",
        "    def _build_reward_fn(\n",
        "            self,\n",
        "            timeframe_size: int,\n",
        "            target_horizon_len: int,\n",
        "            highs: np.ndarray,\n",
        "            lows: np.ndarray,\n",
        "            closes: np.ndarray\n",
        "    ) -> np.ndarray:\n",
        "        return np.float32([[\n",
        "            np.log(closes[i: i + target_horizon_len].max()/closes[i - 1]),\n",
        "            np.log(closes[i - 1]/closes[i: i + target_horizon_len].min())\n",
        "        ] for i in range(timeframe_size, closes.shape[0] - target_horizon_len + 1)])\n"
      ],
      "metadata": {
        "id": "mMikRRe3aTzA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# environments.rewards.marketlimitorder import MarketLimitOrderRF"
      ],
      "metadata": {
        "id": "eduAxGxth02d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from environments.rewards.function import RewardFunction\n",
        "\n",
        "\n",
        "class MarketLimitOrderRF(RewardFunction):\n",
        "    def __init__(\n",
        "            self,\n",
        "            timeframe_size: int,\n",
        "            target_horizon_len: int,\n",
        "            highs: np.ndarray,\n",
        "            lows: np.ndarray,\n",
        "            closes: np.ndarray,\n",
        "            fees_percentage: float\n",
        "    ):\n",
        "        super().__init__(\n",
        "            timeframe_size=timeframe_size,\n",
        "            target_horizon_len=target_horizon_len,\n",
        "            highs=highs,\n",
        "            lows=lows,\n",
        "            closes=closes,\n",
        "            fees_percentage=fees_percentage\n",
        "        )\n",
        "\n",
        "    def _build_reward_fn(\n",
        "            self,\n",
        "            timeframe_size: int,\n",
        "            target_horizon_len: int,\n",
        "            highs: np.ndarray,\n",
        "            lows: np.ndarray,\n",
        "            closes: np.ndarray\n",
        "    ) -> np.ndarray:\n",
        "        return np.float32([[\n",
        "            np.log(highs[i: i + target_horizon_len].max()/closes[i - 1]),\n",
        "            np.log(closes[i - 1]/lows[i: i + target_horizon_len].min())\n",
        "        ] for i in range(timeframe_size, closes.shape[0] - target_horizon_len + 1)])\n"
      ],
      "metadata": {
        "id": "BTlazyG6aUTf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# metrics.trading.pnl import CumulativeLogReturn"
      ],
      "metadata": {
        "id": "A_CKByqqivoq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from metrics.metric import Metric\n",
        "\n",
        "\n",
        "class CumulativeLogReturn(Metric):\n",
        "    def __init__(self):\n",
        "        super().__init__(name='Cumulative Log Returns')\n",
        "        self._log_pnl_sum = 0\n",
        "\n",
        "    def reset(self):\n",
        "        self._log_pnl_sum = 0\n",
        "\n",
        "    def update(self, log_pnl: float):\n",
        "        self._log_pnl_sum += log_pnl\n",
        "\n",
        "    def result(self) -> float:\n",
        "        return self._log_pnl_sum\n"
      ],
      "metadata": {
        "id": "A7pa5Ah3dggV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# metrics.trading.risk import InvestmentRisk"
      ],
      "metadata": {
        "id": "D-EixBGSivzj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from metrics.metric import Metric\n",
        "\n",
        "\n",
        "class InvestmentRisk(Metric):\n",
        "    def __init__(self):\n",
        "        super().__init__(name='Investment Risk')\n",
        "        self._sum_good_transactions = 0\n",
        "        self._sum_bad_transactions = 0\n",
        "\n",
        "    def reset(self):\n",
        "        self._sum_good_transactions = 0\n",
        "        self._sum_bad_transactions = 0\n",
        "\n",
        "    def update(self, log_pnl: float):\n",
        "        if log_pnl > 0:\n",
        "            self._sum_good_transactions += 1\n",
        "        elif log_pnl < 0:\n",
        "            self._sum_bad_transactions += 1\n",
        "        else:\n",
        "            return\n",
        "\n",
        "    def result(self) -> float:\n",
        "        total_investments = (self._sum_bad_transactions + self._sum_good_transactions)\n",
        "        return 0 if total_investments == 0 else self._sum_bad_transactions/total_investments\n"
      ],
      "metadata": {
        "id": "D0NPPDeNdg_k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# metrics.trading.sharpe import SharpeRatio"
      ],
      "metadata": {
        "id": "rQhz87Ukiv90"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from metrics.metric import Metric\n",
        "\n",
        "\n",
        "class SharpeRatio(Metric):\n",
        "    def __init__(self):\n",
        "        super().__init__(name='Sharpe')\n",
        "        self._episode_log_pnls = []\n",
        "\n",
        "    def reset(self):\n",
        "        self._episode_log_pnls = []\n",
        "\n",
        "    def update(self, log_pnl: float):\n",
        "        self._episode_log_pnls.append(log_pnl)\n",
        "\n",
        "    def result(self) -> float:\n",
        "        episode_log_returns = np.float64(self._episode_log_pnls)\n",
        "        average_returns = episode_log_returns.mean()\n",
        "        std_returns = episode_log_returns.std()\n",
        "        return np.exp(average_returns/std_returns)\n"
      ],
      "metadata": {
        "id": "jvKQqnPudhZS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# metrics.trading.sortino import SortinoRatio"
      ],
      "metadata": {
        "id": "Kyfm9-WMiwoY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from metrics.metric import Metric\n",
        "\n",
        "\n",
        "class SortinoRatio(Metric):\n",
        "    def __init__(self):\n",
        "        super().__init__(name='Sortino')\n",
        "        self._episode_log_pnls = []\n",
        "\n",
        "    def reset(self):\n",
        "        self._episode_log_pnls = []\n",
        "\n",
        "    def update(self, log_pnl: float):\n",
        "        self._episode_log_pnls.append(log_pnl)\n",
        "\n",
        "    def result(self) -> float:\n",
        "        episode_log_returns = np.float64(self._episode_log_pnls)\n",
        "        average_returns = episode_log_returns.mean()\n",
        "        std_downfall_returns = episode_log_returns[episode_log_returns < 0].std()\n",
        "        return np.exp(average_returns/std_downfall_returns)\n"
      ],
      "metadata": {
        "id": "GgWvMk5jdh4W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# metrics.trading.drawdown import MaximumDrawdown"
      ],
      "metadata": {
        "id": "-MSMi4ETixD3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from metrics.metric import Metric\n",
        "\n",
        "\n",
        "class MaximumDrawdown(Metric):\n",
        "    def __init__(self):\n",
        "        super().__init__(name='Maximum Drawdown')\n",
        "        self._log_pnl_sum = 0\n",
        "        self._log_pnl_sum_peak = 0\n",
        "        self._hourly_mdds = []\n",
        "\n",
        "    def reset(self):\n",
        "        self._log_pnl_sum = 0\n",
        "        self._log_pnl_sum_peak = 0\n",
        "        self._hourly_mdds = []\n",
        "\n",
        "    def update(self, log_pnl: float):\n",
        "        self._log_pnl_sum += log_pnl\n",
        "\n",
        "        if self._log_pnl_sum_peak < self._log_pnl_sum:\n",
        "            self._log_pnl_sum_peak = self._log_pnl_sum\n",
        "\n",
        "        self._hourly_mdds.append(1 if self._log_pnl_sum_peak == 0 else self._log_pnl_sum/self._log_pnl_sum_peak)\n",
        "\n",
        "    def result(self) -> float:\n",
        "        log_mdd = min(self._hourly_mdds)\n",
        "        return 1 - log_mdd\n"
      ],
      "metadata": {
        "id": "nAGMGciUdipK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# rules.nconsecutive import NConsecutive    NOT YET IMPLEMENTED"
      ],
      "metadata": {
        "id": "oVrSQLZGiqdw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import numpy as np\n",
        "# from environments.actions import Action\n",
        "# from rules.rule import Rule\n",
        "\n",
        "\n",
        "class NConsecutive(Rule):\n",
        "    def __init__(self, window_size: int):\n",
        "        self._window_size = window_size\n",
        "        self._actions_queue = []\n",
        "\n",
        "    def filter(self, action: int) -> int:\n",
        "        if len(self._actions_queue) < self._window_size:\n",
        "            self._actions_queue.insert(0, action)\n",
        "            return Action.HOLD.value\n",
        "\n",
        "        self._actions_queue.pop(-1)\n",
        "        self._actions_queue.insert(0, action)\n",
        "        return action if len(set(self._actions_queue)) == 1 else Action.HOLD.value\n"
      ],
      "metadata": {
        "id": "HIOCxAAFcdrx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training **agents**"
      ],
      "metadata": {
        "id": "d05s4OOiWHgH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# !pip install config"
      ],
      "metadata": {
        "id": "zwm65zpqfF3A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# database.entities.crypto import Crypto"
      ],
      "metadata": {
        "id": "4BbG8Qnkjao9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from database.entities.crypto import Crypto\n",
        "\n",
        "class Crypto:\n",
        "    def __init__(self, symbol: str, name: str, start_year: int):\n",
        "        self._symbol = symbol\n",
        "        self._name = name\n",
        "        self._start_year = start_year\n",
        "\n",
        "    @property\n",
        "    def symbol(self) -> str:\n",
        "        return self._symbol\n",
        "\n",
        "    @property\n",
        "    def name(self) -> str:\n",
        "        return self._name\n",
        "\n",
        "    @property\n",
        "    def start_year(self) -> int:\n",
        "        return self._start_year\n"
      ],
      "metadata": {
        "id": "MYtql8A3jKDu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# database.network.network import DatasetDownloader"
      ],
      "metadata": {
        "id": "cXsqWJa2jw22"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# import pandas as pd\n",
        "# from database.entities.crypto import Crypto\n",
        "# from abc import ABC, abstractmethod\n",
        "\n",
        "\n",
        "class DatasetDownloader(ABC):\n",
        "    def __init__(self, date_column_name: str, verbose: bool):\n",
        "        self._date_column_name = date_column_name\n",
        "        self._verbose = verbose\n",
        "\n",
        "    @property\n",
        "    def date_column_name(self) -> str:\n",
        "        return self._date_column_name\n",
        "\n",
        "    @property\n",
        "    def verbose(self) -> bool:\n",
        "        return self._verbose\n",
        "\n",
        "    def _store_dataset(self, dataset_df: pd.DataFrame, filepath: str, columns: list or None = None):\n",
        "        assert not dataset_df.duplicated(subset=self.date_column_name).any(), \\\n",
        "            f'AssertionError: Date column is expected to be unique, got duplicates'\n",
        "\n",
        "        assert dataset_df[self.date_column_name].is_monotonic_increasing, \\\n",
        "            f'AssertionError: Date column is expected to be monotonic and increasing'\n",
        "\n",
        "        dataset_df.to_csv(filepath, columns=columns, index=False)\n",
        "\n",
        "    @abstractmethod\n",
        "    def download_historical_data(self, crypto: Crypto, history_filepath: str) -> bool:\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def update_historical_data(self, crypto: Crypto, history_filepath: str) -> bool:\n",
        "        pass\n"
      ],
      "metadata": {
        "id": "jAS0eqJIkJFL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# database.network.coinapi.coinapi import CoinAPIDownloader\n"
      ],
      "metadata": {
        "id": "qCwX4wOkkRIO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "from urllib.parse import urlencode"
      ],
      "metadata": {
        "id": "h35QGU3Aj1ed"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import requests\n",
        "# from abc import ABC, abstractmethod\n",
        "# from urllib.parse import urlencode\n",
        "# from database.network.network import DatasetDownloader\n",
        "\n",
        "\n",
        "class CoinAPIDownloader(DatasetDownloader, ABC):\n",
        "    def __init__(self, verbose: bool):\n",
        "        super().__init__(date_column_name=self._get_date_column_name(), verbose=verbose)\n",
        "\n",
        "        self._api_key_list = [\n",
        "            '70E10174-E29D-449F-9F2E-6E8362931DD9',\n",
        "            '27E5E40C-7A6B-45EB-A5C8-8311B049A741',\n",
        "            '8F6252DE-0AD7-478F-91C7-141141E8BE8B',\n",
        "            '3B49210E-100B-4F8D-9011-2BA5D38274BA',\n",
        "            'BF6BF46F-B44B-416E-9656-2D2AAFBC058B',\n",
        "            'B21A98A2-C953-4C73-84CF-CFFB6F712200',\n",
        "            '51667E99-7686-4496-B23D-6DA54F7E37AE',\n",
        "            '0921F87B-BF55-4B78-B8B0-E023B4D7A2E2',\n",
        "            '3F9E3251-029C-457A-9ADA-7F21A440AAF9',\n",
        "            '41EBEA2D-1A4B-4654-8A41-186639B9AB9F',\n",
        "            '6B93AEC2-910C-4064-80FB-91AED487AB97',\n",
        "            '83049379-23DE-4CB0-8299-7137BB836D48',\n",
        "            'B08FCA1F-F454-4C34-AC01-42F16354BCBC',\n",
        "            '12E5D72C-25A6-4ED6-8384-7C291EC43768',\n",
        "            '4F287859-5A00-47EF-AC91-8A2629F8C6A1',\n",
        "            '3744F705-2C4A-406C-AA96-EB1B557A84EF',\n",
        "            '3F77D500-457E-4A96-9CE1-1DEF3FC7033B',\n",
        "            '455C2228-0D6F-4B62-8336-4BAA24C1A46E',\n",
        "            '7E37E058-670C-4ED6-B7BE-DC00F309D9FF',\n",
        "            '0F517C3D-162C-4C5E-AE18-544B201C9BC0'\n",
        "        ]\n",
        "\n",
        "    @property\n",
        "    def api_key_list(self) -> list[str]:\n",
        "        return self._api_key_list\n",
        "\n",
        "    @abstractmethod\n",
        "    def _get_date_column_name(self) -> str:\n",
        "        pass\n",
        "\n",
        "    @abstractmethod\n",
        "    def _get_request_params(self) -> dict[str, str]:\n",
        "        pass\n",
        "\n",
        "    @staticmethod\n",
        "    def _encode_request_url(base_url: str, request_params: dict, api_key: str) -> str:\n",
        "        request_params['apikey'] = api_key\n",
        "        encoded_params = urlencode(request_params)\n",
        "        return f'{base_url}?{encoded_params}'\n",
        "\n",
        "    def _get_response(self, base_url: str, request_params: dict) -> requests.Response or None:\n",
        "        for api_key in self._api_key_list:\n",
        "            if self._verbose:\n",
        "                print(f'Using apikey: {api_key}')\n",
        "\n",
        "            encoded_request_url = self._encode_request_url(\n",
        "                base_url=base_url,\n",
        "                request_params=request_params,\n",
        "                api_key=api_key\n",
        "            )\n",
        "            response = requests.get(encoded_request_url)\n",
        "\n",
        "            if self._verbose:\n",
        "                print(f'Response Status: {response.status_code} - {response.reason}')\n",
        "\n",
        "            if response.status_code == 200:\n",
        "                return response\n",
        "        return None\n"
      ],
      "metadata": {
        "id": "W5st5Ba1jxAH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# database.network.coinapi.ohlcv import OHLCVDownloader"
      ],
      "metadata": {
        "id": "djcyC-6zjfil"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import io"
      ],
      "metadata": {
        "id": "1zaWZHjHjTWv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import io\n",
        "# import pandas as pd\n",
        "# from enum import Enum\n",
        "# from database.entities.crypto import Crypto\n",
        "# from database.network.coinapi.coinapi import CoinAPIDownloader\n",
        "\n",
        "\n",
        "class OHLCVDownloader(CoinAPIDownloader):\n",
        "    class HistoricalFrequency(Enum):\n",
        "        MINUTE = '1MIN'\n",
        "        HOUR = '1HRS'\n",
        "\n",
        "    def __init__(self, historical_frequency: HistoricalFrequency or str, verbose: bool):\n",
        "        if isinstance(historical_frequency, str):\n",
        "            if historical_frequency == '1HRS':\n",
        "                self._historical_frequency = self.HistoricalFrequency.HOUR\n",
        "            elif historical_frequency == '1MIN':\n",
        "                self._historical_frequency = self.HistoricalFrequency.MINUTE\n",
        "            else:\n",
        "                raise NotImplementedError(f'\"{historical_frequency}\" frequency has not been implemented yet')\n",
        "        else:\n",
        "            self._historical_frequency = historical_frequency\n",
        "\n",
        "        super().__init__(verbose=verbose)\n",
        "\n",
        "        self._history_request_url = 'https://rest.coinapi.io/v1/ohlcv/{}/USD/history'\n",
        "        self._latest_request_url = 'https://rest.coinapi.io/v1/ohlcv/{}/USD/latest'\n",
        "        self._download_limit = 100000\n",
        "        self._update_limit = 1000\n",
        "\n",
        "    def _get_date_column_name(self) -> str:\n",
        "        return 'time_period_end'\n",
        "\n",
        "    def _get_request_params(self) -> dict:\n",
        "        return {\n",
        "            'period_id': self._historical_frequency.value,\n",
        "            'output_format': 'csv',\n",
        "            'csv_set_delimiter': ',',\n",
        "            'time_start': '{}-{}-{}T00:00:00',\n",
        "            'limit': '{}'\n",
        "        }\n",
        "\n",
        "    def download_historical_data(self, crypto: Crypto, history_filepath: str) -> bool:\n",
        "        if self.verbose:\n",
        "            print(f'Downloading {crypto.name} market history data for {crypto.start_year}')\n",
        "\n",
        "        request_params = self._get_request_params()\n",
        "        request_params['time_start'] = request_params['time_start'].format(crypto.start_year, '01', '01')\n",
        "        request_params['limit'] = request_params['limit'].format(self._download_limit)\n",
        "        base_url = self._history_request_url.format(crypto.symbol)\n",
        "\n",
        "        response = self._get_response(\n",
        "            base_url=base_url,\n",
        "            request_params=request_params\n",
        "        )\n",
        "        if response is not None and response.status_code == 200:\n",
        "            ohlcv_df = pd.read_csv(io.StringIO(response.text), sep=',')\n",
        "            super()._store_dataset(dataset_df=ohlcv_df, filepath=history_filepath)\n",
        "            return True\n",
        "        else:\n",
        "            return False\n",
        "\n",
        "    def update_historical_data(self, crypto: Crypto, history_filepath: str) -> bool:\n",
        "        if self.verbose:\n",
        "            print(f'Updating {crypto.name} market latest data for {crypto.start_year}')\n",
        "\n",
        "        request_params = self._get_request_params()\n",
        "        request_params['limit'] = request_params['limit'].format(self._update_limit)\n",
        "        del request_params['time_start']\n",
        "        base_url = self._latest_request_url.format(crypto.symbol)\n",
        "\n",
        "        response = self._get_response(\n",
        "            base_url=base_url,\n",
        "            request_params=request_params\n",
        "        )\n",
        "\n",
        "        if response is not None and response.status_code == 200:\n",
        "            history_df = pd.read_csv(history_filepath)\n",
        "            latest_df = pd.read_csv(io.StringIO(response.text), sep=',').sort_values(\n",
        "                by=self.date_column_name, ascending=True\n",
        "            )\n",
        "            merged_df = pd.concat((history_df, latest_df), ignore_index=True)\n",
        "            merged_df.drop_duplicates(subset=self.date_column_name, inplace=True)\n",
        "\n",
        "            super()._store_dataset(dataset_df=merged_df, filepath=history_filepath)\n",
        "            return True\n",
        "        return False\n"
      ],
      "metadata": {
        "id": "4HcW-zACk7at"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# config - import config.py"
      ],
      "metadata": {
        "id": "Szr-E6wnkurc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# from database.entities.crypto import Crypto\n",
        "# from database.network.coinapi.ohlcv import OHLCVDownloader\n",
        "\n",
        "class Config:\n",
        "\n",
        "    # --- Database ---\n",
        "    supported_cryptos = {\n",
        "        'BTC': Crypto(symbol='BTC', name='bitcoin', start_year=2017),\n",
        "        'ETH': Crypto(symbol='ETH', name='ethereum', start_year=2017),\n",
        "        'SOL': Crypto(symbol='SOL',  name='solana', start_year=2020),\n",
        "        'ADA': Crypto(symbol='ADA', name='ada', start_year=2017),\n",
        "        'BNB': Crypto(symbol='BNB', name='bnb', start_year=2019),\n",
        "        'XRP': Crypto(symbol='XRP', name='xrp', start_year=2019),\n",
        "        'DOGE': Crypto(symbol='DOGE', name='doge', start_year=2020),\n",
        "        'MATIC': Crypto(symbol='MATIC', name='polygon', start_year=2020),\n",
        "        'TRON': Crypto(symbol='TRON', name='tron', start_year=2018),\n",
        "        'LTC': Crypto(symbol='LTC', name='litecoin', start_year=2018),\n",
        "        'DOT': Crypto(symbol='DOT', name='polkadot', start_year=2021),\n",
        "        'AVAX': Crypto(symbol='AVAX', name='avalanche', start_year=2021),\n",
        "        'XMR': Crypto(symbol='XMR', name='monero', start_year=2018),\n",
        "        'BAT': Crypto(symbol='BAT', name='basic authentication token', start_year=2018),\n",
        "        'LRC': Crypto(symbol='LRC', name='loopring', start_year=2018)\n",
        "    }\n",
        "\n",
        "    ohlcv_dataset_period_id = OHLCVDownloader.HistoricalFrequency.HOUR\n",
        "    ohlcv_history_filepath = 'database/storage/downloads/ohlcv/{}.csv'\n",
        "    gtrends_history_filepath = 'database/storage/downloads/gtrends/{}.csv'\n",
        "    dataset_save_filepath = '/content/drive/MyDrive/Colab Notebooks/Sample Tradernet/Binance_BTCUSDT_d_EDIT.csv'\n",
        "    # dataset_save_filepath = 'database/storage/datasets/{}.csv'\n",
        "    all_features = [\n",
        "        'date', 'open', 'high', 'low', 'close', 'volume', 'trades',\n",
        "        'open_log_returns', 'high_log_returns', 'low_log_returns',\n",
        "        'close_log_returns', 'volume_log_returns', 'trades_log_returns', 'hour',\n",
        "        'dema', 'vwap', 'bband_up', 'bband_down', 'adl', 'obv',\n",
        "        'macd_signal_diffs', 'stoch', 'aroon_up', 'aroon_down', 'rsi', 'adx', 'cci',\n",
        "        'close_dema', 'close_vwap', 'bband_up_close', 'close_bband_down', 'adl_diffs2', 'obv_diffs2', 'trends'\n",
        "    ]\n",
        "    regression_features = [\n",
        "        'open_log_returns', 'high_log_returns', 'low_log_returns',\n",
        "        'close_log_returns', 'volume_log_returns', 'trades_log_returns', 'hour',\n",
        "        'macd_signal_diffs', 'stoch', 'aroon_up', 'aroon_down', 'rsi', 'adx', 'cci',\n",
        "        'close_dema', 'close_vwap', 'bband_up_close', 'close_bband_down', 'adl_diffs2', 'obv_diffs2', 'trends'\n",
        "    ]\n",
        "\n",
        "    # --- Model ---\n",
        "    checkpoint_dir = 'database/storage/checkpoints/'\n",
        "\n",
        "    # --- Clustering ---\n",
        "    crypto_clusters = [\n",
        "        ['BTC', 'ETH', 'SOL', 'ADA', 'XPR', 'DOGE', 'DOT', 'AVAX', 'BAT', 'LRC'],\n",
        "        ['ETH', 'BNB', 'MATIC', 'TRON', 'LTC', 'XMR']\n",
        "    ]\n"
      ],
      "metadata": {
        "id": "ur6Dcav9jKt_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# import config\n",
        "from sklearn.preprocessing import MinMaxScaler"
      ],
      "metadata": {
        "id": "j1RmYzcEWHsi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Reading Datasets for Tradernet\n",
        "\n",
        "\n",
        "def read_dataset_for_tradernet(\n",
        "        dataset_filepath,\n",
        "        timeframe_size,\n",
        "        target_horizon_len,\n",
        "        num_eval_samples,\n",
        "        fees,\n",
        "        reward_fn_instance\n",
        "):\n",
        "    # Reading dataset\n",
        "    crypto_dataset_df = pd.read_csv(Config.dataset_save_filepath.format(dataset_filepath))\n",
        "    samples_df = crypto_dataset_df[Config.regression_features]\n",
        "\n",
        "    # Scaling data\n",
        "    scaler = MinMaxScaler(feature_range=(0, 1.0))\n",
        "    samples = samples_df.to_numpy(dtype=np.float32)\n",
        "\n",
        "    num_train_scale_samples = samples.shape[0] - num_eval_samples - target_horizon_len - timeframe_size + 1\n",
        "    samples[: num_train_scale_samples] = scaler.fit_transform(samples[: num_train_scale_samples])\n",
        "    samples[num_train_scale_samples: ] = scaler.transform(samples[num_train_scale_samples: ])\n",
        "\n",
        "    # Constructing timeframes for train, test\n",
        "    inputs = np.float32([samples[i: i + timeframe_size] for i in range(samples.shape[0] - timeframe_size - target_horizon_len + 1)])\n",
        "\n",
        "    # Splitting inputs to train-test data\n",
        "    num_train_inputs = inputs.shape[0] - num_eval_samples\n",
        "    x_train = inputs[: num_train_inputs]\n",
        "    x_eval = inputs[num_train_inputs:]\n",
        "\n",
        "    # Computing reward functions for train, test data\n",
        "    closes = crypto_dataset_df['close'].to_numpy(dtype=np.float32)\n",
        "    highs = crypto_dataset_df['high'].to_numpy(dtype=np.float32)\n",
        "    lows = crypto_dataset_df['low'].to_numpy(dtype=np.float32)\n",
        "\n",
        "    train_reward_fn = reward_fn_instance(\n",
        "        timeframe_size=timeframe_size,\n",
        "        target_horizon_len=target_horizon_len,\n",
        "        highs=highs[: samples.shape[0] - num_eval_samples],\n",
        "        lows=lows[: samples.shape[0] - num_eval_samples],\n",
        "        closes=closes[: samples.shape[0] - num_eval_samples],\n",
        "        fees_percentage=fees\n",
        "    )\n",
        "\n",
        "    eval_reward_fn = reward_fn_instance(\n",
        "        timeframe_size=timeframe_size,\n",
        "        target_horizon_len=target_horizon_len,\n",
        "        highs=highs[samples.shape[0] - num_eval_samples - timeframe_size - target_horizon_len + 1:],\n",
        "        lows=lows[samples.shape[0] - num_eval_samples - timeframe_size - target_horizon_len + 1:],\n",
        "        closes=closes[samples.shape[0] - num_eval_samples - timeframe_size - target_horizon_len + 1:],\n",
        "        fees_percentage=fees\n",
        "    )\n",
        "\n",
        "    assert x_train.shape[0] == train_reward_fn.get_reward_fn_shape()[0], \\\n",
        "        f'AssertionError: DimensionMismatch: x_train: {x_train.shape}, train_reward_fn: {train_reward_fn.get_reward_fn_shape()}'\n",
        "    assert x_eval.shape[0] == eval_reward_fn.get_reward_fn_shape()[0], \\\n",
        "        f'AssertionError: DimensionMismatch: x_eval: {x_eval.shape}, eval_reward_fn: {eval_reward_fn.get_reward_fn_shape()}'\n",
        "\n",
        "    return x_train, train_reward_fn, x_eval, eval_reward_fn"
      ],
      "metadata": {
        "id": "eYDU2YyVcTf2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Reading Datasets for Smurf Agent\n",
        "\n",
        "\n",
        "def read_dataset_for_smurf_agent(  #### CHANGE VARIABLE\n",
        "        dataset_filepath,\n",
        "        timeframe_size,\n",
        "        target_horizon_len,\n",
        "        num_eval_samples,\n",
        "        fees,\n",
        "        reward_fn_instance\n",
        "):\n",
        "    # Reading dataset\n",
        "    crypto_dataset_df = pd.read_csv(Config.dataset_save_filepath.format(dataset_filepath))     #### CHANGE VARIABLE\n",
        "    samples_df = crypto_dataset_df[Config.regression_features]        #### CHANGE VARIABLE\n",
        "\n",
        "    # Scaling data\n",
        "    scaler = MinMaxScaler(feature_range=(0, 1.0))         #### CHANGE VARIABLE\n",
        "    samples = samples_df.to_numpy(dtype=np.float32)         #### CHANGE VARIABLE\n",
        "\n",
        "    num_train_scale_samples = samples.shape[0] - num_eval_samples - target_horizon_len - timeframe_size + 1         #### CHANGE VARIABLE\n",
        "    samples[: num_train_scale_samples] = scaler.fit_transform(samples[: num_train_scale_samples])           #### CHANGE VARIABLE\n",
        "    samples[num_train_scale_samples: ] = scaler.transform(samples[num_train_scale_samples: ])         #### CHANGE VARIABLE\n",
        "\n",
        "    # Constructing timeframes for train, test\n",
        "    inputs = np.float32([samples[i: i + timeframe_size] for i in range(samples.shape[0] - timeframe_size - target_horizon_len + 1)])      #### CHANGE VARIABLE\n",
        "\n",
        "    # Splitting inputs to train-test data\n",
        "    num_train_inputs = inputs.shape[0] - num_eval_samples         #### CHANGE VARIABLE\n",
        "    x_train = inputs[: num_train_inputs]              #### CHANGE VARIABLE\n",
        "    x_eval = inputs[num_train_inputs:]                #### CHANGE VARIABLE\n",
        "\n",
        "    # Computing reward functions for train, test data\n",
        "    closes = crypto_dataset_df['close'].to_numpy(dtype=np.float32)      #### CHANGE VARIABLE\n",
        "    highs = crypto_dataset_df['high'].to_numpy(dtype=np.float32)        #### CHANGE VARIABLE\n",
        "    lows = crypto_dataset_df['low'].to_numpy(dtype=np.float32)          #### CHANGE VARIABLE\n",
        "\n",
        "    train_reward_fn = SmurfRewardFunction(reward_function=reward_fn_instance(       #### CHANGE VARIABLE\n",
        "        timeframe_size=timeframe_size,\n",
        "        target_horizon_len=target_horizon_len,\n",
        "        highs=highs[: samples.shape[0] - num_eval_samples],\n",
        "        lows=lows[: samples.shape[0] - num_eval_samples],\n",
        "        closes=closes[: samples.shape[0] - num_eval_samples],\n",
        "        fees_percentage=fees\n",
        "    ))\n",
        "\n",
        "    eval_reward_fn = SmurfRewardFunction(reward_function=reward_fn_instance(        #### CHANGE VARIABLE\n",
        "        timeframe_size=timeframe_size,\n",
        "        target_horizon_len=target_horizon_len,\n",
        "        highs=highs[samples.shape[0] - num_eval_samples - timeframe_size - target_horizon_len + 1:],\n",
        "        lows=lows[samples.shape[0] - num_eval_samples - timeframe_size - target_horizon_len + 1:],\n",
        "        closes=closes[samples.shape[0] - num_eval_samples - timeframe_size - target_horizon_len + 1:],\n",
        "        fees_percentage=fees\n",
        "    ))\n",
        "\n",
        "    assert x_train.shape[0] == train_reward_fn.get_reward_fn_shape()[0], \\\n",
        "        f'AssertionError: DimensionMismatch: x_train: {x_train.shape}, train_reward_fn: {train_reward_fn.get_reward_fn_shape()}'\n",
        "    assert x_eval.shape[0] == eval_reward_fn.get_reward_fn_shape()[0], \\\n",
        "        f'AssertionError: DimensionMismatch: x_eval: {x_eval.shape}, eval_reward_fn: {eval_reward_fn.get_reward_fn_shape()}'\n",
        "\n",
        "    return x_train, train_reward_fn, x_eval, eval_reward_fn                          #### CHANGE VARIABLE"
      ],
      "metadata": {
        "id": "tyWVzcu3W3Q4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Building Tradernet Agent\n",
        "\n",
        "\n",
        "def build_tradernet_agent(            #### CHANGE VARIABLE\n",
        "        agent_instance,\n",
        "        observation_spec,\n",
        "        action_spec,\n",
        "        time_step_spec,\n",
        "        env_batch_size,\n",
        "        checkpoint_filepath,\n",
        "        fc_layers,\n",
        "        conv_layers\n",
        "):\n",
        "    return agent_instance(            #### CHANGE VARIABLE\n",
        "        input_tensor_spec=observation_spec,\n",
        "        action_spec=action_spec,\n",
        "        time_step_spec=time_step_spec,\n",
        "        env_batch_size=env_batch_size,\n",
        "        checkpoint_filepath=checkpoint_filepath,\n",
        "        fc_layers=fc_layers,\n",
        "        conv_layers=conv_layers\n",
        "    )"
      ],
      "metadata": {
        "id": "GsKDCmOKdE5j"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Building Smurf Agent\n",
        "\n",
        "\n",
        "def build_smurf_agent(            #### CHANGE VARIABLE\n",
        "        agent_instance,\n",
        "        observation_spec,\n",
        "        action_spec,\n",
        "        time_step_spec,\n",
        "        env_batch_size,\n",
        "        checkpoint_filepath,\n",
        "        fc_layers,\n",
        "        conv_layers\n",
        "):\n",
        "    return agent_instance(        #### CHANGE VARIABLE\n",
        "        input_tensor_spec=observation_spec,\n",
        "        action_spec=action_spec,\n",
        "        time_step_spec=time_step_spec,\n",
        "        env_batch_size=env_batch_size,\n",
        "        checkpoint_filepath=checkpoint_filepath,\n",
        "        fc_layers=fc_layers,\n",
        "        conv_layers=conv_layers\n",
        "    )"
      ],
      "metadata": {
        "id": "tVQHrSN4XEKd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Building Tradernet Agent Trainer\n",
        "\n",
        "\n",
        "def train_tradernet_agent(\n",
        "        dataset_filepath,\n",
        "        timeframe_size,\n",
        "        target_horizon_len,\n",
        "        num_eval_samples,\n",
        "        fees,\n",
        "        reward_fn_instance,\n",
        "        agent_instance,\n",
        "        checkpoint_filepath,\n",
        "        fc_layers,\n",
        "        conv_layers,\n",
        "        train_episode_steps,\n",
        "        train_iterations,\n",
        "        eval_episodes,\n",
        "        steps_per_eval,\n",
        "        steps_per_log,\n",
        "        steps_per_checkpoint,\n",
        "        save_best_only\n",
        "):\n",
        "    x_train, train_reward_fn, x_eval, eval_reward_fn = read_dataset_for_tradernet(        #### CHANGED\n",
        "        dataset_filepath=dataset_filepath,\n",
        "        timeframe_size=timeframe_size,\n",
        "        target_horizon_len=target_horizon_len,\n",
        "        num_eval_samples=num_eval_samples,\n",
        "        fees=fees,\n",
        "        reward_fn_instance=reward_fn_instance\n",
        "    )\n",
        "    train_env = TFTradingEnvironment(\n",
        "        env=TradingEnvironment(env_config={\n",
        "            'states': x_train,\n",
        "            'reward_fn': train_reward_fn,\n",
        "            'episode_steps': train_episode_steps,\n",
        "            'metrics': [CumulativeLogReturn(), InvestmentRisk(), SharpeRatio(), SortinoRatio(), MaximumDrawdown()]\n",
        "        })\n",
        "    )\n",
        "    eval_env = TFTradingEnvironment(\n",
        "        env=TradingEnvironment(env_config={\n",
        "            'states': x_eval,\n",
        "            'reward_fn': eval_reward_fn,\n",
        "            'episode_steps': x_eval.shape[0] - 1,\n",
        "            'metrics': [CumulativeLogReturn(), InvestmentRisk(), SharpeRatio(), SortinoRatio(), MaximumDrawdown()]\n",
        "        })\n",
        "    )\n",
        "\n",
        "    tf_train_env = TFPyEnvironment(environment=train_env)\n",
        "    tf_eval_env = TFPyEnvironment(environment=eval_env)\n",
        "\n",
        "    agent = build_tradernet_agent(        #### CHANGE VARIABLE\n",
        "        agent_instance=agent_instance,\n",
        "        observation_spec=tf_train_env.observation_spec(),\n",
        "        action_spec=tf_train_env.action_spec(),\n",
        "        time_step_spec=tf_train_env.time_step_spec(),\n",
        "        env_batch_size=tf_train_env.batch_size,\n",
        "        checkpoint_filepath=checkpoint_filepath,\n",
        "        fc_layers=fc_layers,\n",
        "        conv_layers=conv_layers,\n",
        "    )\n",
        "\n",
        "    agent.initialize()\n",
        "\n",
        "    eval_avg_returns = agent.train(       #### CHANGE VARIABLE\n",
        "        train_env=tf_train_env,\n",
        "        eval_env=tf_eval_env,\n",
        "        train_iterations=train_iterations,\n",
        "        eval_episodes=eval_episodes,\n",
        "        iterations_per_eval=steps_per_eval,\n",
        "        iterations_per_log=steps_per_log,\n",
        "        iterations_per_checkpoint=steps_per_checkpoint,\n",
        "        save_best_only=save_best_only\n",
        "    )\n",
        "    eval_metrics = eval_env.get_metrics()\n",
        "    return eval_avg_returns, eval_metrics"
      ],
      "metadata": {
        "id": "Ht0trWhOdW5O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Building Smurf Agent Trainer\n",
        "\n",
        "\n",
        "def train_smurf_agent(\n",
        "        dataset_filepath,\n",
        "        timeframe_size,\n",
        "        target_horizon_len,\n",
        "        num_eval_samples,\n",
        "        fees,\n",
        "        reward_fn_instance,\n",
        "        agent_instance,\n",
        "        checkpoint_filepath,\n",
        "        fc_layers,\n",
        "        conv_layers,\n",
        "        train_episode_steps,\n",
        "        train_iterations,\n",
        "        eval_episodes,\n",
        "        steps_per_eval,\n",
        "        steps_per_log,\n",
        "        steps_per_checkpoint,\n",
        "        save_best_only\n",
        "):\n",
        "    x_train, train_reward_fn, x_eval, eval_reward_fn = read_dataset_for_smurf_agent(      #### CHANGED\n",
        "        dataset_filepath=dataset_filepath,\n",
        "        timeframe_size=timeframe_size,\n",
        "        target_horizon_len=target_horizon_len,\n",
        "        num_eval_samples=num_eval_samples,\n",
        "        fees=fees,\n",
        "        reward_fn_instance=reward_fn_instance\n",
        "    )\n",
        "    train_env = TFTradingEnvironment(      #### CHANGE VARIABLE\n",
        "        env=TradingEnvironment(env_config={\n",
        "            'states': x_train,\n",
        "            'reward_fn': train_reward_fn,\n",
        "            'episode_steps': train_episode_steps,\n",
        "            'metrics': [CumulativeLogReturn(), InvestmentRisk(), SharpeRatio(), SortinoRatio(), MaximumDrawdown()]\n",
        "        })\n",
        "    )\n",
        "    eval_env = TFTradingEnvironment(       #### CHANGE VARIABLE\n",
        "        env=TradingEnvironment(env_config={\n",
        "            'states': x_eval,\n",
        "            'reward_fn': eval_reward_fn,\n",
        "            'episode_steps': x_eval.shape[0] - 1,\n",
        "            'metrics': [CumulativeLogReturn(), InvestmentRisk(), SharpeRatio(), SortinoRatio(), MaximumDrawdown()]\n",
        "        })\n",
        "    )\n",
        "\n",
        "    tf_train_env = TFPyEnvironment(environment=train_env)          #### CHANGE VARIABLE\n",
        "    tf_eval_env = TFPyEnvironment(environment=eval_env)              #### CHANGE VARIABLE\n",
        "\n",
        "    agent = build_smurf_agent(                   #### CHANGE VARIABLE\n",
        "        agent_instance=agent_instance,\n",
        "        observation_spec=tf_train_env.observation_spec(),\n",
        "        action_spec=tf_train_env.action_spec(),\n",
        "        time_step_spec=tf_train_env.time_step_spec(),\n",
        "        env_batch_size=tf_train_env.batch_size,\n",
        "        checkpoint_filepath=checkpoint_filepath,\n",
        "        fc_layers=fc_layers,\n",
        "        conv_layers=conv_layers,\n",
        "    )\n",
        "\n",
        "    agent.initialize()\n",
        "\n",
        "    eval_avg_returns = agent.train(            #### CHANGE VARIABLE\n",
        "        train_env=tf_train_env,\n",
        "        eval_env=tf_eval_env,\n",
        "        train_iterations=train_iterations,\n",
        "        eval_episodes=eval_episodes,\n",
        "        iterations_per_eval=steps_per_eval,\n",
        "        iterations_per_log=steps_per_log,\n",
        "        iterations_per_checkpoint=steps_per_checkpoint,\n",
        "        save_best_only=save_best_only\n",
        "    )\n",
        "    eval_metrics = eval_env.get_metrics()          #### CHANGE VARIABLE\n",
        "    return eval_avg_returns, eval_metrics"
      ],
      "metadata": {
        "id": "0FHqstYEXIlu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Building Train Configs for Tradernet Agent\n",
        "\n",
        "tradernet_datasets_dict = {'BTC': 'BTC', 'ETH': 'ETH', 'ADA': 'ADA', 'XRP': 'XRP', 'LTC': 'LTC'}\n",
        "tradernet_rewards_dict = {\n",
        "    'Market-Orders':  MarketOrderRF,\n",
        "    'Market-Limit Orders': MarketLimitOrderRF\n",
        "}\n",
        "tradernet_agents_configs = {\n",
        "    'DDQN': {\n",
        "        'agent_instance': DQNAgent,\n",
        "        'train_iterations': 50000,\n",
        "        'steps_per_eval': 500,\n",
        "        'steps_per_log': 500,\n",
        "        'steps_per_checkpoint': 500\n",
        "    },\n",
        "    'PPO': {\n",
        "        'agent_instance': PPOAgent,\n",
        "        'train_iterations': 1000,\n",
        "        'steps_per_eval': 10,\n",
        "        'steps_per_log': 10,\n",
        "        'steps_per_checkpoint': 10\n",
        "    }\n",
        "}\n",
        "tradernet_train_dict = {\n",
        "    'timeframe_size': 12,\n",
        "    'target_horizon_len': 20,\n",
        "    'num_eval_samples': 2250,\n",
        "    'fees': 0.007,\n",
        "    'fc_layers': [256, 256],\n",
        "    'conv_layers': [(32, 3, 1)],\n",
        "    'train_episode_steps': 100,\n",
        "    'eval_episodes': 1,\n",
        "    'save_best_only': True\n",
        "}"
      ],
      "metadata": {
        "id": "kqFJKsZGd4cA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Building Train Configs for Smurf Agent\n",
        "\n",
        "smurf_datasets_dict = {'BTC': 'BTC', 'ETH': 'ETH', 'ADA': 'ADA', 'XRP': 'XRP', 'LTC': 'LTC'}\n",
        "smurf_rewards_dict = {\n",
        "    'Market-Orders':  MarketOrderRF,\n",
        "    'Market-Limit Orders': MarketLimitOrderRF\n",
        "}\n",
        "smurf_agents_configs = {\n",
        "    'PPO': {\n",
        "        'agent_instance': PPOAgent,\n",
        "        'train_iterations': 1000,\n",
        "        'steps_per_eval': 10,\n",
        "        'steps_per_log': 10,\n",
        "        'steps_per_checkpoint': 10\n",
        "    },\n",
        "    'DDQN': {\n",
        "        'agent_instance': DQNAgent,\n",
        "        'train_iterations': 50000,\n",
        "        'steps_per_eval': 500,\n",
        "        'steps_per_log': 500,\n",
        "        'steps_per_checkpoint': 500\n",
        "    }\n",
        "}\n",
        "smurf_train_dict = {\n",
        "    'timeframe_size': 12,\n",
        "    'target_horizon_len': 20,\n",
        "    'num_eval_samples': 2250,\n",
        "    'fees': 0.01,\n",
        "    'fc_layers': [256, 256],\n",
        "    'conv_layers': [(32, 3, 1)],\n",
        "    'train_episode_steps': 100,\n",
        "    'eval_episodes': 1,\n",
        "    'save_best_only': True\n",
        "}"
      ],
      "metadata": {
        "id": "gxEpRNdWXQnH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Run TraderNet Experiments\n",
        "\n",
        "tradernet_results = {\n",
        "    'PPO': {dataset_name: {} for dataset_name in tradernet_datasets_dict.keys()},\n",
        "    'DDQN': {dataset_name: {} for dataset_name in tradernet_datasets_dict.keys()}\n",
        "}\n",
        "\n",
        "tradernet_colors = {\n",
        "    'BTC': 'green',\n",
        "    'ETH': 'blue',\n",
        "    'XRP': 'red',\n",
        "    'ADA': 'black',\n",
        "    'LTC': 'orange'\n",
        "}\n",
        "tradernet_linestyles = {\n",
        "    'Market-Orders': '--',\n",
        "    'Market-Limit-Orders': '-'\n",
        "}\n",
        "\n",
        "for agent_name, agent_config in tradernet_agents_configs.items():\n",
        "    for dataset_name, dataset_filepath in tradernet_datasets_dict.items():\n",
        "        for reward_fn_name, reward_fn_instance in tradernet_rewards_dict.items():\n",
        "            tf.random.set_seed(seed=0)\n",
        "\n",
        "            train_params = {\n",
        "                'dataset_filepath': dataset_filepath,\n",
        "                'reward_fn_instance': reward_fn_instance,\n",
        "                'checkpoint_filepath': f'database/storage/checkpoints/experiments/tradernet/{agent_name}/{dataset_name}/{reward_fn_name}/',\n",
        "                **tradernet_train_dict,\n",
        "                **agent_config\n",
        "            }\n",
        "            eval_avg_returns, eval_metrics = train_tradernet_agent(**train_params)\n",
        "\n",
        "            tradernet_results[agent_name][dataset_name][reward_fn_name] = (eval_avg_returns, eval_metrics)\n",
        "\n",
        "        for reward_fn_name, reward_fn_results in tradernet_results[agent_name][dataset_name].items():\n",
        "            eval_avg_returns, eval_metrics = reward_fn_results\n",
        "\n",
        "            metrics_dict = {\n",
        "                'steps': [10000*i for i in range(len(eval_avg_returns))],\n",
        "                'average_returns': eval_avg_returns,\n",
        "                **{metric.name: metric.episode_metrics for metric in eval_metrics}\n",
        "            }\n",
        "            metrics_df = pd.DataFrame(metrics_dict)\n",
        "            metrics_df.to_csv(f'experiments/tradernet/{agent_name}/{dataset_name}_{reward_fn_name}.csv', index=False)\n"
      ],
      "metadata": {
        "id": "zcXc0iyReZNt",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 425
        },
        "outputId": "bdc1c93a-22ab-4b46-a421-a868bd6f59b7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "KeyError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-70-50d194ead447>\u001b[0m in \u001b[0;36m<cell line: 20>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     30\u001b[0m                 \u001b[0;34m**\u001b[0m\u001b[0magent_config\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     31\u001b[0m             }\n\u001b[0;32m---> 32\u001b[0;31m             \u001b[0meval_avg_returns\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0meval_metrics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_tradernet_agent\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0mtrain_params\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     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m             \u001b[0mtradernet_results\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0magent_name\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdataset_name\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mreward_fn_name\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0meval_avg_returns\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0meval_metrics\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-66-0bdaafb97f85>\u001b[0m in \u001b[0;36mtrain_tradernet_agent\u001b[0;34m(dataset_filepath, timeframe_size, target_horizon_len, num_eval_samples, fees, reward_fn_instance, agent_instance, checkpoint_filepath, fc_layers, conv_layers, train_episode_steps, train_iterations, eval_episodes, steps_per_eval, steps_per_log, steps_per_checkpoint, save_best_only)\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0msave_best_only\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m ):\n\u001b[0;32m---> 23\u001b[0;31m     x_train, train_reward_fn, x_eval, eval_reward_fn = read_dataset_for_tradernet(        #### CHANGED\n\u001b[0m\u001b[1;32m     24\u001b[0m         \u001b[0mdataset_filepath\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdataset_filepath\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0mtimeframe_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtimeframe_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-62-82d6f3f85901>\u001b[0m in \u001b[0;36mread_dataset_for_tradernet\u001b[0;34m(dataset_filepath, timeframe_size, target_horizon_len, num_eval_samples, fees, reward_fn_instance)\u001b[0m\n\u001b[1;32m     12\u001b[0m     \u001b[0;31m# Reading dataset\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0mcrypto_dataset_df\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mConfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset_save_filepath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdataset_filepath\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---> 14\u001b[0;31m     \u001b[0msamples_df\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcrypto_dataset_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mConfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mregression_features\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     15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0;31m# Scaling data\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/core/frame.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m   3811\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mis_iterator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3812\u001b[0m                 \u001b[0mkey\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3813\u001b[0;31m             \u001b[0mindexer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcolumns\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_indexer_strict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"columns\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\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   3814\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3815\u001b[0m         \u001b[0;31m# take() does not accept boolean indexers\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/core/indexes/base.py\u001b[0m in \u001b[0;36m_get_indexer_strict\u001b[0;34m(self, key, axis_name)\u001b[0m\n\u001b[1;32m   6068\u001b[0m             \u001b[0mkeyarr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindexer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnew_indexer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reindex_non_unique\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkeyarr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   6069\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 6070\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_raise_if_missing\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkeyarr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindexer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis_name\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   6071\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   6072\u001b[0m         \u001b[0mkeyarr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtake\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindexer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/pandas/core/indexes/base.py\u001b[0m in \u001b[0;36m_raise_if_missing\u001b[0;34m(self, key, indexer, axis_name)\u001b[0m\n\u001b[1;32m   6128\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0muse_interval_msg\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   6129\u001b[0m                     \u001b[0mkey\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 6130\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mKeyError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"None of [{key}] are in the [{axis_name}]\"\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   6131\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   6132\u001b[0m             \u001b[0mnot_found\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mensure_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmissing_mask\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnonzero\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munique\u001b[0m\u001b[0;34m(\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;31mKeyError\u001b[0m: \"None of [Index(['open_log_returns', 'high_log_returns', 'low_log_returns',\\n       'close_log_returns', 'volume_log_returns', 'trades_log_returns', 'hour',\\n       'macd_signal_diffs', 'stoch', 'aroon_up', 'aroon_down', 'rsi', 'adx',\\n       'cci', 'close_dema', 'close_vwap', 'bband_up_close', 'close_bband_down',\\n       'adl_diffs2', 'obv_diffs2', 'trends'],\\n      dtype='object')] are in the [columns]\""
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Run TraderNet Experiments with Smurf\n",
        "\n",
        "smurf_results = {\n",
        "    'PPO': {dataset_name: {} for dataset_name in smurf_datasets_dict.keys()},\n",
        "    'DDQN': {dataset_name: {} for dataset_name in smurf_datasets_dict.keys()}\n",
        "}\n",
        "\n",
        "smurf_colors = {\n",
        "    'BTC': 'green',\n",
        "    'ETH': 'blue',\n",
        "    'XRP': 'red',\n",
        "    'ADA': 'black',\n",
        "    'LTC': 'orange'\n",
        "}\n",
        "smurf_linestyles = {\n",
        "    'Market-Orders': '--',\n",
        "    'Market-Limit-Orders': '-'\n",
        "}\n",
        "\n",
        "for agent_name, agent_config in smurf_agents_configs.items():\n",
        "    for dataset_name, dataset_filepath in smurf_datasets_dict.items():\n",
        "        for reward_fn_name, reward_fn_instance in smurf_rewards_dict.items():\n",
        "            tf.random.set_seed(seed=0)\n",
        "\n",
        "            train_params = {\n",
        "                'dataset_filepath': dataset_filepath,\n",
        "                'reward_fn_instance': reward_fn_instance,\n",
        "                'checkpoint_filepath': f'database/storage/checkpoints/experiments/smurf/{agent_name}/{dataset_name}/{reward_fn_name}/',\n",
        "                **smurf_train_dict,\n",
        "                **agent_config\n",
        "            }\n",
        "            eval_avg_returns, eval_metrics = train_smurf_agent(**train_params)\n",
        "\n",
        "            smurf_results[agent_name][dataset_name][reward_fn_name] = (eval_avg_returns, eval_metrics)\n",
        "\n",
        "        for reward_fn_name, reward_fn_results in smurf_results[agent_name][dataset_name].items():\n",
        "            eval_avg_returns, eval_metrics = reward_fn_results\n",
        "\n",
        "            metrics_dict = {\n",
        "                'steps': [10000*i for i in range(len(eval_avg_returns))],\n",
        "                'average_returns': eval_avg_returns,\n",
        "                **{metric.name: metric.episode_metrics for metric in eval_metrics}\n",
        "            }\n",
        "            metrics_df = pd.DataFrame(metrics_dict)\n",
        "            metrics_df.to_csv(f'experiments/smurf/{agent_name}/{dataset_name}_{reward_fn_name}.csv', index=False)"
      ],
      "metadata": {
        "id": "dCrNqtjtXVit"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}