ai_week_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "ai_week.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "ongEoQjaTJQk",
        "G2__PdzbNrPi",
        "RBxwrlSzHZLu",
        "tO6bmhG1HZL0",
        "mCwoQwnGHZL5",
        "wNelbGV7HZME",
        "3Eqej3v50a91"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ongEoQjaTJQk",
        "colab_type": "text"
      },
      "source": [
        "# Intstall coach\n",
        "Just use pip"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sRS7Z_lfGrs6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pip install rl_coach"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9qn7Qf9gHZLk"
      },
      "source": [
        "# AI Week Workshop "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kEcvEIxqHZLt"
      },
      "source": [
        "### ***Add new environmen***t\n",
        "\n",
        "In this section we will implement the short corridor environment from Sutton & Barto Book.\n",
        "\n",
        "*   Three non terminal states- The location of the agent\n",
        "\n",
        "*   The observations are one-hot encoding of the states\n",
        "*   Actions are reversed in the second state\n",
        "\n",
        "\n",
        "*   Reward is -1 for each time step\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G2__PdzbNrPi",
        "colab_type": "text"
      },
      "source": [
        "##### ***Helper function*** \n",
        "The following code snippet contains some defines and an one-hot encoding helper function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lnMACvXo_y3U",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%writefile short_corridor_env_helpper.py\n",
        "import numpy as np\n",
        "\n",
        "LEFT = 0\n",
        "RIGHT = 1\n",
        "GOAL_STATE = 3\n",
        "NUM_STATES = 4\n",
        "REVERSE_STATE = 1\n",
        "\n",
        "def to_one_hot(state):\n",
        "    observation = np.zeros((NUM_STATES,))\n",
        "    observation[state] = 1\n",
        "    return observation"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RBxwrlSzHZLu"
      },
      "source": [
        "##### ***Write short corridor environment*** \n",
        "Compete the following functions:\n",
        " function and the step function\n",
        "\n",
        "1.   is_done - will return a boolean . True only at termination state\n",
        "\n",
        "2.   reset - Resets environment to initial state\n",
        "3.   step - Returns the next observation, reward, and the boolean flag done\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "* **complete code**\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "6f517868-0681-4d69-e5a0-2bd2c8a78bd1",
        "id": "uxwi5S1vHZLw",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile short_corridor_env.py\n",
        "import numpy as np\n",
        "import gym\n",
        "from gym import spaces\n",
        "from  short_corridor_env_helpper import *\n",
        "\n",
        "\n",
        "class ShortCorridorEnv(gym.Env):\n",
        "\n",
        "    def __init__(self):\n",
        "        # Class constructor- Initializes class variables and sets initial state\n",
        "        self.observation_space = spaces.Box(0, 1, shape=(NUM_STATES,))\n",
        "        self.action_space = spaces.Discrete(2)\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        '''\n",
        "        Resets the environment to start state\n",
        "        '''\n",
        "        # Boolean. True only if the goal state is reached\n",
        "        self.goal_reached = ???\n",
        "        # An integer representing the state. Number between zero and three\n",
        "        self.current_state = ???\n",
        "        observation = to_one_hot(???)\n",
        "        return observation\n",
        "\n",
        "    def _is_done(self, current_state):\n",
        "        '''\n",
        "        return done a Boolean- True only if we reached the goal state\n",
        "        '''\n",
        "        ???\n",
        "        return done\n",
        "\n",
        "    def step(self, action):\n",
        "        '''\n",
        "        Returns the next observation, reward, and the boolean flag done\n",
        "        '''\n",
        "\n",
        "        if action ==LEFT:\n",
        "          step = -1\n",
        "        elif action == RIGHT:\n",
        "          step = ???\n",
        "\n",
        "        if self.current_state == REVERSE_STATE:\n",
        "        ### Replace step = -1 with step = 1 and vise versa\n",
        "            ???\n",
        "\n",
        "        self.current_state += step\n",
        "        self.current_state = max(0, self.current_state)\n",
        "\n",
        "        observation = to_one_hot(self.current_state)\n",
        "        reward = ???\n",
        "        done = self._is_done(self.current_state)\n",
        "\n",
        "        return observation, reward, done, {}\n",
        "\n"
      ],
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting short_corridor_env.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tO6bmhG1HZL0"
      },
      "source": [
        "##### ***Write preset to run existing agent on the new environment***\n",
        "*We will use the same preset from DQN example*.\n",
        "\n",
        "Since our environment is already using Gym API we are almost good to go.\n",
        "\n",
        "When selecting the environment parametes in the preset use **GymEnvironmentParameters** and pass the path of the environment source code using the level parameter"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "wVmJ2CfbHZL1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "da081f19-98bd-45cc-cde6-b8e16fefad66"
      },
      "source": [
        "%%writefile short_corridor_dqn_preset.py\n",
        "from rl_coach.environments.gym_environment import GymEnvironmentParameters\n",
        "from rl_coach.filters.filter import NoInputFilter, NoOutputFilter\n",
        "from rl_coach.agents.dqn_agent import DQNAgentParameters\n",
        "from rl_coach.graph_managers.basic_rl_graph_manager import BasicRLGraphManager\n",
        "from rl_coach.graph_managers.graph_manager import SimpleSchedule\n",
        "from rl_coach.memories.memory import MemoryGranularity\n",
        "\n",
        "\n",
        "####################\n",
        "# Graph Scheduling #\n",
        "####################\n",
        "schedule_params = SimpleSchedule()\n",
        "\n",
        "\n",
        "#########\n",
        "# Agent #\n",
        "#########\n",
        "agent_params = DQNAgentParameters()\n",
        "agent_params.input_filter = NoInputFilter()\n",
        "agent_params.output_filter = NoOutputFilter()\n",
        "# DQN params\n",
        "# ER size\n",
        "agent_params.memory.max_size = (MemoryGranularity.Transitions, 40000)\n",
        "\n",
        "\n",
        "###############\n",
        "# Environment #\n",
        "###############\n",
        "env_params = GymEnvironmentParameters(level='short_corridor_env:ShortCorridorEnv')\n",
        "\n",
        "\n",
        "#################\n",
        "# Graph Manager #\n",
        "#################\n",
        "graph_manager = BasicRLGraphManager(agent_params=agent_params,\n",
        "                                    env_params=env_params,\n",
        "                                    schedule_params=schedule_params)\n"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Writing short_corridor_dqn_preset.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mCwoQwnGHZL5"
      },
      "source": [
        "##### ***Run new preset***"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "4db-w4ETHZL6",
        "colab": {}
      },
      "source": [
        "!coach -p /content/short_corridor_dqn_preset.py:graph_manager\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aMs0JnUBHZL8"
      },
      "source": [
        "### ***Add new agent***\n",
        "Coach modularity makes adding an agent a clean and simple task.\n",
        "Tipicaly consists of four parts:\n",
        "\n",
        "\n",
        "1.   Implement an agent spesific network head (and loss)\n",
        "2.   Implement exploration policy (optional)\n",
        "3.   Define new parametes class that extends `AgentParametes`\n",
        "4.   Implement a preset to run the agent on some environment\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MtfJT8XbmsUo",
        "colab_type": "text"
      },
      "source": [
        "##### ***Write stochastic output layer***\n",
        "We use stochastic policy, meaning that we only produce the probability of going left and going right.\n",
        "This layer takes in the input from previous layer, the middleware, and outputs two numbers. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KIYd8t8GnCSu",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "76bb7100-76d1-4be3-cc55-22e4cc7acbd8"
      },
      "source": [
        "%%writefile probabilistic_layer.py\n",
        "import tensorflow as tf\n",
        "from rl_coach.architectures.tensorflow_components.layers import Dense\n",
        "\n",
        "class ProbabilisticLayer(object):\n",
        "    def __init__(self, input_layer, num_actions):\n",
        "        super().__init__()\n",
        "        scores = Dense(num_actions)(input_layer, name='logit')\n",
        "        self.event_probs = tf.nn.softmax(scores, name=\"policy\")\n",
        "        # define the distributions for the policy and the old policy\n",
        "        self.policy_distribution = tf.contrib.distributions.Categorical(probs=self.event_probs)\n",
        "\n",
        "    def log_prob(self, action):\n",
        "        return self.policy_distribution.log_prob(action)\n",
        "\n",
        "    def layer_output(self):\n",
        "        return self.event_probs"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Writing probabilistic_layer.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "wNelbGV7HZME"
      },
      "source": [
        "##### ***Implement network head i.e. implement the loss***\n",
        "The Head needs to inherit from the base class `Head`.\n",
        "\n",
        "Inorder to maximize the sum of rewards, we want to go in the following direction $-\\Sigma_i A_i \\nabla_Wlog(\\pi(a_i|x_i))$\n",
        "\n",
        "$- A_i \\nabla_Wlog[\\pi(a_i|x_i)]$\n",
        "\n",
        "`Complete code`\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "dd33ee5d-2eba-4e45-a230-009bd4b37fd2",
        "id": "hUnA2ZJHHZMF",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile simple_pg_head.py\n",
        "import tensorflow as tf\n",
        "from rl_coach.architectures.tensorflow_components.heads.head import Head\n",
        "from rl_coach.base_parameters import AgentParameters\n",
        "from rl_coach.spaces import SpacesDefinition\n",
        "from probabilistic_layer import ProbabilisticLayer\n",
        "\n",
        "\n",
        "class SimplePgHead(Head):\n",
        "    def __init__(self, agent_parameters: AgentParameters,\n",
        "                 spaces: SpacesDefinition, network_name: str,\n",
        "                 head_idx: int = 0, is_local: bool = True):\n",
        "        super().__init__(agent_parameters, spaces, network_name)\n",
        "\n",
        "        self.exploration_policy = agent_parameters.exploration\n",
        "\n",
        "    def _build_module(self, input_layer):\n",
        "        # Define inputs\n",
        "        actions = tf.placeholder(tf.int32, [None], name=\"actions\")\n",
        "        advantages = tf.placeholder(tf.float32, [None], name=\"advantages\")\n",
        "\n",
        "        # Two actions, left or right\n",
        "        policy_distribution = ProbabilisticLayer(input_layer, num_actions=2)\n",
        "\n",
        "        # calculate loss\n",
        "        log_prob = policy_distribution.log_prob(???)\n",
        "        mudulated_log_prob = ???\n",
        "        expected_mudulated_log_prob = tf.reduce_mean(mudulated_log_prob)\n",
        "\n",
        "        ### Coach bookkeeping\n",
        "        # List of placeholders for additional inputs to the head \n",
        "        #(except from the middleware input)\n",
        "        self.input.append(???)\n",
        "        # The output of the head, which is also the output of the network.\n",
        "        self.output.append(policy_distribution.layer_output)#???\n",
        "        # Placeholder for the target that we will use to train the network\n",
        "        self.target = ???\n",
        "        # The loss that we will use to train the network\n",
        "        self.loss = ???\n",
        "        tf.losses.add_loss(self.loss)\n",
        "\n"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting simple_pg_head.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "3Eqej3v50a91"
      },
      "source": [
        "##### ***Define exploration policy*** \n",
        "Every iteration we want to sample from the network output distribution i.e. toss a bias coin to get the agent actual move\n",
        "\n",
        "**`Complete code`**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "4efed15d-3956-4ff1-f072-4ec6592da0e0",
        "id": "zuYBPgnI0a-A",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile simple_pg_exploration.py\n",
        "\n",
        "import numpy as np\n",
        "from rl_coach.exploration_policies.exploration_policy import ExplorationPolicy, ExplorationParameters\n",
        "from rl_coach.spaces import ActionSpace\n",
        "\n",
        "\n",
        "class DiscreteExplorationParameters(ExplorationParameters):\n",
        "    @property\n",
        "    def path(self):\n",
        "        return 'simple_pg_exploration:DiscreteExploration'\n",
        "\n",
        "\n",
        "class DiscreteExploration(ExplorationPolicy):\n",
        "    \"\"\"\n",
        "    Discrete exploration policy is intended for discrete action spaces. It expects the action values to\n",
        "    represent a probability distribution over the action\n",
        "    \"\"\"\n",
        "    def __init__(self, action_space: ActionSpace):\n",
        "        \"\"\"\n",
        "        :param action_space: the action space used by the environment\n",
        "        \"\"\"\n",
        "        super().__init__(action_space)\n",
        "\n",
        "    def get_action(self, probabilities):\n",
        "        # choose actions according to the probabilities\n",
        "        action = np.random.choice(self.action_space.actions, p=???)\n",
        "        return chosen_action, probabilities\n"
      ],
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting simple_pg_exploration.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XG5y-G43HZL9"
      },
      "source": [
        "##### ***Define new agent parameters***\n",
        "Coach is modular!\n",
        "\n",
        "Each class in Coach has a complementary parameters class which defines its constructor. \n",
        "This is also true for the agent. The agent has a complementary `AgentParameters` class. This class enable to select the paramenters of the agent sub modules.\n",
        "\n",
        "It consists of the following four parts:\n",
        "\n",
        "\n",
        "\n",
        "1.   algorithm\n",
        "2.   exploration\n",
        "3.   memory\n",
        "4.   Networks\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "047d9045-f825-4456-8340-328db9c78bac",
        "id": "bTAcPbnsHZL-",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile simple_pg_params.py\n",
        "from rl_coach.architectures.embedder_parameters import InputEmbedderParameters\n",
        "from rl_coach.architectures.head_parameters import HeadParameters\n",
        "from rl_coach.architectures.middleware_parameters import FCMiddlewareParameters\n",
        "from rl_coach.base_parameters import NetworkParameters, AlgorithmParameters, \\\n",
        "    AgentParameters\n",
        "\n",
        "from rl_coach.exploration_policies.additive_noise import AdditiveNoiseParameters\n",
        "from rl_coach.exploration_policies.categorical import CategoricalParameters\n",
        "from rl_coach.memories.episodic.single_episode_buffer import SingleEpisodeBufferParameters\n",
        "from rl_coach.spaces import DiscreteActionSpace, BoxActionSpace\n",
        "from rl_coach.agents.policy_optimization_agent import PolicyGradientRescaler\n",
        "from simple_pg_exploration import DiscreteExplorationParameters\n",
        "\n",
        "class SimplePgAgentParameters(AgentParameters):\n",
        "    def __init__(self):\n",
        "        super().__init__(algorithm=SimplePGAlgorithmParameters(),\n",
        "                         #exploration=CategoricalParameters(),\n",
        "                         exploration=DiscreteExplorationParameters(),\n",
        "                         memory=SingleEpisodeBufferParameters(),\n",
        "                         networks={\"main\": SimplePgTopology()})\n",
        "    @property\n",
        "    def path(self):\n",
        "        #return 'simple_pg_agent:SimplePgAgent'\n",
        "        return 'rl_coach.agents.policy_gradients_agent:PolicyGradientsAgent'\n",
        "\n",
        "        \n",
        "    \n",
        "# Since we are adding a new head we need to tell coach the heads path\n",
        "class SimplePgHeadParams(HeadParameters):\n",
        "    def __init__(self):\n",
        "        super().__init__(parameterized_class_name=\"AiWeekHead\")\n",
        "\n",
        "    @property\n",
        "    def path(self):\n",
        "        return 'simple_pg_head:SimplePgHead'\n",
        "\n",
        "\n",
        "class SimplePgTopology(NetworkParameters):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.input_embedders_parameters = {'observation': InputEmbedderParameters()}\n",
        "        self.middleware_parameters = FCMiddlewareParameters()\n",
        "        self.heads_parameters = [SimplePgHeadParams()]\n",
        "\n",
        "\n",
        "class SimplePGAlgorithmParameters(AlgorithmParameters):\n",
        "    \"\"\"\n",
        "    :param num_steps_between_gradient_updates: (int)\n",
        "        The number of steps between calculating gradients for the collected data. In the A3C paper, this parameter is\n",
        "        called t_max. Since this algorithm is on-policy, only the steps collected between each two gradient calculations\n",
        "        are used in the batch.\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        # TOTAL_RETURN\n",
        "        # FUTURE_RETURN\n",
        "        # FUTURE_RETURN_NORMALIZED_BY_EPISODE \n",
        "        # FUTURE_RETURN_NORMALIZED_BY_TIMESTEP\n",
        "        # Q_VALUE\n",
        "        # A_VALUE\n",
        "        # TD_RESIDUAL\n",
        "        # DISCOUNTED_TD_RESIDUAL\n",
        "        # GAE\n",
        "        self.policy_gradient_rescaler = PolicyGradientRescaler.FUTURE_RETURN\n",
        "        self.num_steps_between_gradient_updates = 20000  # this is called t_max in all the papers\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting simple_pg_params.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_8RT_QFdHZMH"
      },
      "source": [
        "##### ***Write preset to run new agent on short corridor***\n",
        "complete code\n",
        "* **complete code**\n",
        "* **Hint: look at DQN preset**\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "ce90a661-5691-4b91-8bf0-cc9af56df38f",
        "id": "D1PLQuNaHZMI",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile short_corridor_new_agent_preset.py\n",
        "from rl_coach.base_parameters import VisualizationParameters\n",
        "from rl_coach.core_types import EnvironmentEpisodes, EnvironmentSteps\n",
        "from rl_coach.environments.gym_environment import GymEnvironmentParameters\n",
        "from rl_coach.filters.filter import NoInputFilter, NoOutputFilter\n",
        "from rl_coach.graph_managers.basic_rl_graph_manager import BasicRLGraphManager\n",
        "from rl_coach.graph_managers.graph_manager import SimpleSchedule\n",
        "from rl_coach.memories.memory import MemoryGranularity\n",
        "from rl_coach.schedules import LinearSchedule\n",
        "from simple_pg_params import SimplePgAgentParameters\n",
        "\n",
        "\n",
        "####################\n",
        "# Graph Scheduling #\n",
        "####################\n",
        "schedule_params = SimpleSchedule()\n",
        "\n",
        "\n",
        "#########\n",
        "# Agent #\n",
        "#########\n",
        "agent_params = ???\n",
        "agent_params.input_filter = NoInputFilter()\n",
        "agent_params.output_filter = NoOutputFilter()\n",
        "\n",
        "\n",
        "###############\n",
        "# Environment #\n",
        "###############\n",
        "env_params = GymEnvironmentParameters(level='short_corridor_env:ShortCorridorEnv')\n",
        "\n",
        "#################\n",
        "# Graph Manager #\n",
        "#################\n",
        "graph_manager = BasicRLGraphManager(agent_params=agent_params,\n",
        "                                    env_params=env_params,\n",
        "                                    schedule_params=schedule_params)\n",
        "\n"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Writing short_corridor_new_agent_preset.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ufc_EAe3HZMK"
      },
      "source": [
        "##### ***Run preset of the new agent on the new environment***\n",
        "\n",
        "**`Complete code`**\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "ad2xpEu8HZML",
        "colab": {}
      },
      "source": [
        "???"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mH01dwya6w7I"
      },
      "source": [
        "# AI Week Workshop solution"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EEB5Sh7s6w7L"
      },
      "source": [
        "### ***Add new environmen***t\n",
        "\n",
        "In this section we will implement the short corridor environment from Sutton & Barto Book.**bold text**\n",
        "\n",
        "*   Three non terminal states- The location of the agent\n",
        "\n",
        "*   The observations are one-hot encoding of the states\n",
        "*   Actions are reversed in the second state\n",
        "\n",
        "\n",
        "*   Reward is -1 for each time step\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uqJIFrdS6w7M"
      },
      "source": [
        "##### ***Helper function*** \n",
        "The following code snippet contains some defines and an one-hot encoding helper function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "LSVIPyIl6w7M",
        "colab": {}
      },
      "source": [
        "%%writefile short_corridor_env_helpper.py\n",
        "import numpy as np\n",
        "\n",
        "LEFT = 0\n",
        "RIGHT = 1\n",
        "GOAL_STATE = 3\n",
        "NUM_STATES = 4\n",
        "REVERSE_STATE = 1\n",
        "\n",
        "def to_one_hot(state):\n",
        "    observation = np.zeros((NUM_STATES,))\n",
        "    observation[state] = 1\n",
        "    return observation"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "icikXNlP6w7Q"
      },
      "source": [
        "##### ***Write short corridor environment*** \n",
        "Compete the following functions:\n",
        " function and the step function\n",
        "\n",
        "1.   is_done - will return a boolean . True only at termination state\n",
        "\n",
        "2.   reset - Resets environment to initial state\n",
        "3.   step - Returns the next observation, reward, and the boolean flag done\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "* **complete code**\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "6f517868-0681-4d69-e5a0-2bd2c8a78bd1",
        "id": "KcNXvD_x6w7R",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile short_corridor_env.py\n",
        "import numpy as np\n",
        "import gym\n",
        "from gym import spaces\n",
        "from  short_corridor_env_helpper import *#to_one_hot\n",
        "\n",
        "\n",
        "class ShortCorridorEnv(gym.Env):\n",
        "\n",
        "    def __init__(self):\n",
        "        # Class constructor- Initializes class variables and sets initial state\n",
        "        self.observation_space = spaces.Box(0, 1, shape=(NUM_STATES,))\n",
        "        self.action_space = spaces.Discrete(2)\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        '''\n",
        "        Resets the environment to start state\n",
        "        '''\n",
        "        # Boolean. True only if the goal state is reached\n",
        "        self.goal_reached = False#???\n",
        "        # An integer representing the state. Number between zero and three\n",
        "        self.current_state = 0#???\n",
        "        observation = to_one_hot(self.current_state)#???\n",
        "        return observation\n",
        "\n",
        "    def _is_done(self, current_state):\n",
        "        '''\n",
        "        return done a Boolean- True only if we reached the goal state\n",
        "        '''\n",
        "        done = (current_state == GOAL_STATE)#???\n",
        "        return done\n",
        "\n",
        "    def step(self, action):\n",
        "        '''\n",
        "        Returns the next observation, reward, and the boolean flag done\n",
        "        '''\n",
        "\n",
        "        if action ==LEFT:\n",
        "          step = -1\n",
        "        elif action == RIGHT:\n",
        "          step = 1 #???\n",
        "\n",
        "        if self.current_state == REVERSE_STATE:\n",
        "        ### Replace step = -1 with step = 1 and vise versa\n",
        "            step = -step#???\n",
        "\n",
        "        self.current_state += step\n",
        "        self.current_state = max(0, self.current_state)\n",
        "\n",
        "        observation = to_one_hot(self.current_state)\n",
        "        reward = -1#???\n",
        "        done = self._is_done(self.current_state)\n",
        "\n",
        "        return observation, reward, done, {}\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting short_corridor_env.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8IHUjH3n6w7V"
      },
      "source": [
        "##### ***Write preset to run existing agent on the new environment***\n",
        "*We will use the same preset from DQN example*.\n",
        "\n",
        "Since our environment is already using Gym API we are almost good to go.\n",
        "\n",
        "When selecting the environment parametes in the preset use **GymEnvironmentParameters** and pass the path of the environment source code using the level parameter"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "da081f19-98bd-45cc-cde6-b8e16fefad66",
        "id": "xzWdrnB26w7V",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile short_corridor_dqn_preset.py\n",
        "from rl_coach.environments.gym_environment import GymEnvironmentParameters\n",
        "from rl_coach.filters.filter import NoInputFilter, NoOutputFilter\n",
        "from rl_coach.agents.dqn_agent import DQNAgentParameters\n",
        "from rl_coach.graph_managers.basic_rl_graph_manager import BasicRLGraphManager\n",
        "from rl_coach.graph_managers.graph_manager import SimpleSchedule\n",
        "from rl_coach.memories.memory import MemoryGranularity\n",
        "\n",
        "\n",
        "####################\n",
        "# Graph Scheduling #\n",
        "####################\n",
        "schedule_params = SimpleSchedule()\n",
        "\n",
        "\n",
        "#########\n",
        "# Agent #\n",
        "#########\n",
        "agent_params = DQNAgentParameters()\n",
        "agent_params.input_filter = NoInputFilter()\n",
        "agent_params.output_filter = NoOutputFilter()\n",
        "# DQN params\n",
        "# ER size\n",
        "agent_params.memory.max_size = (MemoryGranularity.Transitions, 40000)\n",
        "\n",
        "\n",
        "###############\n",
        "# Environment #\n",
        "###############\n",
        "env_params = GymEnvironmentParameters(level='short_corridor_env:ShortCorridorEnv')\n",
        "\n",
        "\n",
        "#################\n",
        "# Graph Manager #\n",
        "#################\n",
        "graph_manager = BasicRLGraphManager(agent_params=agent_params,\n",
        "                                    env_params=env_params,\n",
        "                                    schedule_params=schedule_params)\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Writing short_corridor_dqn_preset.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "a21At4Gw6w7Y"
      },
      "source": [
        "##### ***Run new preset***"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "7yFMWjXU6w7Z",
        "colab": {}
      },
      "source": [
        "!coach -p /content/short_corridor_dqn_preset.py:graph_manager\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PfKp9Y7G6w7b"
      },
      "source": [
        "### ***Add new agent***"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "N6PvVW456w7c",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NvOwaHaE6w7e"
      },
      "source": [
        "##### ***Write stochastic output layer***\n",
        "We use stochastic policy, meaning that we only produce the probability of going left and going right"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "76bb7100-76d1-4be3-cc55-22e4cc7acbd8",
        "id": "7JmWGTvQ6w7g",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile probabilistic_layer.py\n",
        "import tensorflow as tf\n",
        "from rl_coach.architectures.tensorflow_components.layers import Dense\n",
        "\n",
        "class ProbabilisticLayer(object):\n",
        "    def __init__(self, input_layer, num_actions):\n",
        "        super().__init__()\n",
        "        scores = Dense(num_actions)(input_layer, name='logit')\n",
        "        self.event_probs = tf.nn.softmax(scores, name=\"policy\")\n",
        "        # define the distributions for the policy and the old policy\n",
        "        self.policy_distribution = tf.contrib.distributions.Categorical(probs=self.event_probs)\n",
        "\n",
        "    def log_prob(self, action):\n",
        "        return self.policy_distribution.log_prob(action)\n",
        "\n",
        "    def layer_output(self):\n",
        "        return self.event_probs"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Writing probabilistic_layer.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Q1G7OOnw6w7i"
      },
      "source": [
        "##### ***Define exploration policy*** *italicized text*\n",
        "Every iteration we want to sample from the network output distribution i.e. toss a bias coin to get the agent actual move\n",
        "\n",
        "**`Complete code`**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "4efed15d-3956-4ff1-f072-4ec6592da0e0",
        "id": "5jGmF_T_6w7j",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile simple_pg_exploration.py\n",
        "\n",
        "import numpy as np\n",
        "from rl_coach.exploration_policies.exploration_policy import ExplorationPolicy, ExplorationParameters\n",
        "from rl_coach.spaces import ActionSpace\n",
        "\n",
        "\n",
        "class DiscreteExplorationParameters(ExplorationParameters):\n",
        "    @property\n",
        "    def path(self):\n",
        "        return 'simple_pg_exploration:DiscreteExploration'\n",
        "\n",
        "\n",
        "class DiscreteExploration(ExplorationPolicy):\n",
        "    \"\"\"\n",
        "    Discrete exploration policy is intended for discrete action spaces. It expects the action values to\n",
        "    represent a probability distribution over the action\n",
        "    \"\"\"\n",
        "    def __init__(self, action_space: ActionSpace):\n",
        "        \"\"\"\n",
        "        :param action_space: the action space used by the environment\n",
        "        \"\"\"\n",
        "        super().__init__(action_space)\n",
        "\n",
        "    def get_action(self, probabilities):\n",
        "        # choose actions according to the probabilities\n",
        "        action = np.random.choice(self.action_space.actions,p=probabilities)#p=???)\n",
        "        return chosen_action, probabilities\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting simple_pg_exploration.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "h6AeI9446w7l"
      },
      "source": [
        "##### ***Implement network head i.e. implement the loss***\n",
        "`Complete code`\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "dd33ee5d-2eba-4e45-a230-009bd4b37fd2",
        "id": "yww7sTHy6w7l",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "\n",
        "%%writefile simple_pg_head.py\n",
        "import tensorflow as tf\n",
        "from rl_coach.architectures.tensorflow_components.heads.head import Head\n",
        "from rl_coach.base_parameters import AgentParameters\n",
        "from rl_coach.spaces import SpacesDefinition\n",
        "from probabilistic_layer import ProbabilisticLayer\n",
        "\n",
        "\n",
        "class SimplePgHead(Head):\n",
        "    def __init__(self, agent_parameters: AgentParameters,\n",
        "                 spaces: SpacesDefinition, network_name: str,\n",
        "                 head_idx: int = 0, is_local: bool = True):\n",
        "        super().__init__(agent_parameters, spaces, network_name)\n",
        "\n",
        "        self.exploration_policy = agent_parameters.exploration\n",
        "\n",
        "    def _build_module(self, input_layer):\n",
        "        # Define inputs\n",
        "        actions = tf.placeholder(tf.int32, [None], name=\"actions\")\n",
        "        advantages = tf.placeholder(tf.float32, [None], name=\"advantages\")\n",
        "\n",
        "        # Two actions, left or right\n",
        "        policy_distribution = ProbabilisticLayer(input_layer, num_actions=2)\n",
        "\n",
        "        # calculate loss\n",
        "        log_prob = policy_distribution.log_prob(actions)#actions\n",
        "        mudulated_log_prob = advantages*log_prob\n",
        "        expected_mudulated_log_prob = tf.reduce_mean(mudulated_log_prob)\n",
        "\n",
        "        ### Coach bookkeeping\n",
        "        # List of placeholders for additional inputs to the head \n",
        "        #(except from the middleware input)\n",
        "        self.input.append(actions)#???\n",
        "        # The output of the head, which is also the output of the network.\n",
        "        self.output.append(policy_distribution.layer_output)#???\n",
        "        # Placeholder for the target that we will use to train the network\n",
        "        self.target = advantages#???\n",
        "        # The loss that we will use to train the network\n",
        "        self.loss = -expected_mudulated_log_prob#???\n",
        "        tf.losses.add_loss(self.loss)\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Overwriting simple_pg_head.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "FDbTzz5W6w7o"
      },
      "source": [
        "##### Define new agent parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "5aLSpUPU6w7p",
        "colab": {}
      },
      "source": [
        "%%writefile simple_pg_params.py\n",
        "from rl_coach.architectures.embedder_parameters import InputEmbedderParameters\n",
        "from rl_coach.architectures.head_parameters import HeadParameters\n",
        "from rl_coach.architectures.middleware_parameters import FCMiddlewareParameters\n",
        "from rl_coach.base_parameters import NetworkParameters, AlgorithmParameters, \\\n",
        "    AgentParameters\n",
        "\n",
        "from rl_coach.exploration_policies.additive_noise import AdditiveNoiseParameters\n",
        "from rl_coach.exploration_policies.categorical import CategoricalParameters\n",
        "from rl_coach.memories.episodic.single_episode_buffer import SingleEpisodeBufferParameters\n",
        "from rl_coach.spaces import DiscreteActionSpace, BoxActionSpace\n",
        "from rl_coach.agents.policy_optimization_agent import PolicyGradientRescaler\n",
        "from simple_pg_exploration import DiscreteExplorationParameters\n",
        "\n",
        "class SimplePgAgentParameters(AgentParameters):\n",
        "    def __init__(self):\n",
        "        super().__init__(algorithm=SimplePGAlgorithmParameters(),\n",
        "                         #exploration=CategoricalParameters(),\n",
        "                         exploration=DiscreteExplorationParameters(),\n",
        "                         memory=SingleEpisodeBufferParameters(),\n",
        "                         networks={\"main\": SimplePgTopology()})\n",
        "    @property\n",
        "    def path(self):\n",
        "        #return 'simple_pg_agent:SimplePgAgent'\n",
        "        return 'rl_coach.agents.policy_gradients_agent:PolicyGradientsAgent'\n",
        "\n",
        "        \n",
        "    \n",
        "# Since we are adding a new head we need to tell coach the heads path\n",
        "class SimplePgHeadParams(HeadParameters):\n",
        "    def __init__(self):\n",
        "        super().__init__(parameterized_class_name=\"AiWeekHead\")\n",
        "\n",
        "    @property\n",
        "    def path(self):\n",
        "        return 'simple_pg_head:SimplePgHead'\n",
        "\n",
        "\n",
        "class SimplePgTopology(NetworkParameters):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.input_embedders_parameters = {'observation': InputEmbedderParameters()}\n",
        "        self.middleware_parameters = FCMiddlewareParameters()\n",
        "        self.heads_parameters = [SimplePgHeadParams()]\n",
        "\n",
        "\n",
        "class SimplePGAlgorithmParameters(AlgorithmParameters):\n",
        "    \"\"\"\n",
        "    :param num_steps_between_gradient_updates: (int)\n",
        "        The number of steps between calculating gradients for the collected data. In the A3C paper, this parameter is\n",
        "        called t_max. Since this algorithm is on-policy, only the steps collected between each two gradient calculations\n",
        "        are used in the batch.\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        # TOTAL_RETURN\n",
        "        # FUTURE_RETURN\n",
        "        self.policy_gradient_rescaler = PolicyGradientRescaler.FUTURE_RETURN\n",
        "        self.num_steps_between_gradient_updates = 20000  # this is called t_max in all the papers\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "GrLnoSZm6w7s"
      },
      "source": [
        "##### Write preset to run new agent on short corridor\n",
        "complete code\n",
        "* **complete code**\n",
        "* **Hint: look at DQN preset**\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "outputId": "ce90a661-5691-4b91-8bf0-cc9af56df38f",
        "id": "OZBbiPFo6w7s",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "%%writefile short_corridor_new_agent_preset.py\n",
        "from rl_coach.base_parameters import VisualizationParameters\n",
        "from rl_coach.core_types import EnvironmentEpisodes, EnvironmentSteps\n",
        "from rl_coach.environments.gym_environment import GymEnvironmentParameters\n",
        "from rl_coach.filters.filter import NoInputFilter, NoOutputFilter\n",
        "from rl_coach.graph_managers.basic_rl_graph_manager import BasicRLGraphManager\n",
        "from rl_coach.graph_managers.graph_manager import SimpleSchedule\n",
        "from rl_coach.memories.memory import MemoryGranularity\n",
        "from rl_coach.schedules import LinearSchedule\n",
        "from simple_pg_params import SimplePgAgentParameters\n",
        "\n",
        "\n",
        "####################\n",
        "# Graph Scheduling #\n",
        "####################\n",
        "schedule_params = SimpleSchedule()\n",
        "\n",
        "\n",
        "#########\n",
        "# Agent #\n",
        "#########\n",
        "agent_params = SimplePgAgentParameters()#???\n",
        "agent_params.input_filter = NoInputFilter()\n",
        "agent_params.output_filter = NoOutputFilter()\n",
        "\n",
        "\n",
        "###############\n",
        "# Environment #\n",
        "###############\n",
        "env_params = GymEnvironmentParameters(level='short_corridor_env:ShortCorridorEnv')\n",
        "\n",
        "#################\n",
        "# Graph Manager #\n",
        "#################\n",
        "graph_manager = BasicRLGraphManager(agent_params=agent_params,\n",
        "                                    env_params=env_params,\n",
        "                                    schedule_params=schedule_params)\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Writing short_corridor_new_agent_preset.py\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dxeRBX1g6w7u"
      },
      "source": [
        "##### Run preset of the new agent on the new environment \n",
        "* **Complete code**\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "WgUPBI3P6w7v",
        "colab": {}
      },
      "source": [
        "!coach -p /content/short_corridor_new_agent_preset.py:graph_manager\n",
        "#???"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}