{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DE961vnDvyyq"
      },
      "source": [
        "# Challenge: Multi-Armed Bandits"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8syOlEwVXS8z"
      },
      "source": [
        "## Introduction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5KTU154TdpaT"
      },
      "source": [
        "In this challenge you have two goals.\n",
        "\n",
        "First, you should get familiar with working in notebooks. It is not hard, but it requires some time. Additionaly, there are some features, that will help you to complete the tasks: `solution`s and `check`s. Functions with these parts in them will show you solution and check your solution respectively. `check`s will also give you a part of the key, if you complete the task successfuly.\n",
        "\n",
        "Second, your main goal, is to compare different approaches to solving the multi-armed bandit problem.\n",
        "\n",
        "You will:\n",
        "- Write your own implementations for algorithms;\n",
        "- See how they work both in stationary and in dynamic environments.\n",
        "\n",
        "**NOTE: to complete the tasks, replace `None` and `pass` parts of the code with appropriate values.**\n",
        "\n",
        "**IMPORTANT: when it comes to generation of random numbers, follow the instructions in the tasks.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RyJi86uldbIM"
      },
      "source": [
        "## Dependencies\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L-9R4hB9dr0c"
      },
      "source": [
        "Make sure to install and import all dependencies before you begin to work with the code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "TdTNUq78tfHL"
      },
      "outputs": [],
      "source": [
        "!pip install numpy matplotlib gymnasium git+https://github.com/Codefinity-ML/codefinity-rl.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NVwL1wwekv_c"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import gymnasium as gym\n",
        "from codefinityrl.challenges.utils import np_random_generator\n",
        "from codefinityrl.challenges.mab import *"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YpMq_JgMzIJ3"
      },
      "source": [
        "## Environmets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bfrY_k6gzMZ3"
      },
      "source": [
        "For this challenge, you will use two environments:\n",
        "1. **MultiArmedBanditStationary**: a stationary version of multi-armed bandit problem. All action values are generated at the start, and they remain static until the end;\n",
        "2. **MultiArmedBanditDynamic**: a dynamic version of multi-armed bandit problem. Action values are generated at the start, but they constantly change, shifting linearly from one set of values to another in a certain amount of time steps.\n",
        "\n",
        "Your first task is to create both of these environments using the `gym.make` function. You already know that `gym.make` takes the environment `id` as a parameter. It can also use `max_episode_steps` to set a maximal length of an episode. Additionaly, you can configure the environment by providing environment-specific parameters.\n",
        "\n",
        "These are the parameters for **MultiArmedBanditStationary**:\n",
        "- `id`: \"codefinityrl:MultiArmedBanditStationary-v0\"\n",
        "- `max_episode_steps`: a maximal length of an episode, positive integer or -1 for unlimited length\n",
        "- `n_arms`: a number of arms, positive integer\n",
        "\n",
        "These are the parameters for **MultiArmedBanditDynamic**:\n",
        "- `id`: \"codefinityrl:MultiArmedBanditDynamic-v0\"\n",
        "- `max_episode_steps`: a maximal length of an episode, positive integer or -1 for unlimited length\n",
        "- `n_arms`: a number of arms, positive integer\n",
        "- `drift_interval`: a number of time steps, positive integer. Smaller numbers lead to faster changes in action values\n",
        "\n",
        "Create both environments and choose any values you want for parameters that are listed above(except for ids, values for which are already provided)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "on-2n8uF1xU9"
      },
      "outputs": [],
      "source": [
        "# This variable should contain stationary environment\n",
        "env_stationary = None\n",
        "\n",
        "# This variable should contain dynamic environment\n",
        "env_dynamic = None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ElOu-r95zaGo"
      },
      "outputs": [],
      "source": [
        "# Run this cell and see the solution\n",
        "solution1()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1rC0U13Szplt"
      },
      "outputs": [],
      "source": [
        "# Run this cell and check your solution\n",
        "check1(env_stationary, env_dynamic)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XLiNpn_CJsxj"
      },
      "source": [
        "Now that you have tried writing some code and understood how the `hint`s, `solution`s and `check`s work, you are ready for the real tasks."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_1Jq_JRkfY46"
      },
      "source": [
        "## Epsilon-Greedy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W62Wi2n-0wrN"
      },
      "source": [
        "In this task, you have to write an epsilon-greedy agent.\n",
        "\n",
        "The agent has these parameters:\n",
        "- `n_arms`: a number of arms, positive integer\n",
        "- `epsilon`: a constant that balances eploration and exploitation, float in range from 0 to 1\n",
        "- `alpha`: a rate of change for action values estimates, \"mean\" if estimates are computed as mean of all previous estimates, otherwise - float in range from 0 to 1\n",
        "- `optimistic`: an optimistic estimate of starting action values, non-negative float\n",
        "- `fixed_rng`: a parameter used to verify the correctness of your solution, boolean\n",
        "\n",
        "Your agent should support changes in every parameter listed above, except `fixed_rng`. If you need to generate random values, use `self.np_random` instead of `np.random`.\n",
        "\n",
        "A small reminder:\n",
        "- epsilon-greedy agent selects a random action with probability `epsilon`, otherwise it selects the expected best action;\n",
        "- update formula for action values estimates:\n",
        "   - $Q_t=Q_{t-1}+\\frac{1}{N_i}(R_t-Q_{t-1})$ for mean of rewards;\n",
        "   - $Q_t=Q_{t-1}+\\alpha(R_t-Q_{t-1})$ for constant step size updates."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "djx38Z6L0as2"
      },
      "outputs": [],
      "source": [
        "class EpsilonGreedyAgent:\n",
        "    def __init__(\n",
        "        self,\n",
        "        n_arms: int,\n",
        "        epsilon: float,\n",
        "        alpha: float | str = \"mean\",\n",
        "        optimistic: float = 0,\n",
        "        fixed_rng: bool = True,\n",
        "    ):\n",
        "        self.np_random = np_random_generator(fixed_rng)\n",
        "\n",
        "        # Initialize agent's attributes\n",
        "        self.n_arms = None\n",
        "        self.epsilon = None\n",
        "        self.alpha = None\n",
        "        self.Q = None # Action values\n",
        "        if alpha == \"mean\":\n",
        "            self.N = None # Action selection counts\n",
        "\n",
        "    def select_action(self):\n",
        "        # Write an epsilon-greedy action selection algorithm\n",
        "        if None:\n",
        "            return None # Random action\n",
        "        else:\n",
        "            return None # Current optimal action\n",
        "\n",
        "    def update(self, action: int, reward: float):\n",
        "        # Write an update algorithm\n",
        "        if None:\n",
        "            self.N[action] += None # Update action selection count\n",
        "            self.Q[action] += None # Update action value (mean)\n",
        "        else:\n",
        "            self.Q[action] += None # Update action value (alpha)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "igREZMOwgNRP"
      },
      "outputs": [],
      "source": [
        "solution2()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1dKfON-ngO1v"
      },
      "outputs": [],
      "source": [
        "check2(EpsilonGreedyAgent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0f7a8O6HgI-N"
      },
      "source": [
        "## Upper Confidence Bound"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6KY0hTzJYLS_"
      },
      "source": [
        "In this task, you have to write a UCB agent.\n",
        "\n",
        "The agent has these parameters:\n",
        "- `n_arms`: a number of arms, positive integer\n",
        "- `confidence`: a constant that balances eploration and exploitation, non-negative float\n",
        "\n",
        "Your agent should support changes in every parameter listed above.\n",
        "\n",
        "A small reminder:\n",
        "- before using the formula, UCB should select each action once, otherwise the denominatior in a formula would be 0\n",
        "- the formula for choosing an action looks like this: $A_t = \\underset{a}{\\arg\\max}(Q_t(a)+c\\sqrt\\frac{\\ln(t)}{N_t(a)})$\n",
        "- update formula for action values estimates: $Q_t=Q_{t-1}+\\frac{1}{N_i}(R_t-Q_{t-1})$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X8B-BCypYK_b"
      },
      "outputs": [],
      "source": [
        "class UpperConfidenceBoundAgent:\n",
        "    def __init__(self, n_arms: int, confidence: float):\n",
        "        # Initialize agent's attributes\n",
        "        self.n_arms = None\n",
        "        self.confidence = None\n",
        "        self.Q = None # Action values\n",
        "        self.N = None # Action selection counts\n",
        "        self.t = None # Current time step\n",
        "\n",
        "    def select_action(self):\n",
        "        # Write an action selection algorithm\n",
        "        self.t += None # Increase time step\n",
        "\n",
        "        # Select each action at least once\n",
        "        for action in None:\n",
        "            if None:\n",
        "                return action\n",
        "\n",
        "        return None # The action according to UCB formula\n",
        "\n",
        "    def update(self, action: int, reward: float):\n",
        "        # Write an update algorithm\n",
        "        self.N[action] += None # Update action selection count\n",
        "        self.Q[action] += None # Update action value (mean)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rL8JcvCrVzUV"
      },
      "outputs": [],
      "source": [
        "solution3()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rqo1bpccV02i"
      },
      "outputs": [],
      "source": [
        "check3(UpperConfidenceBoundAgent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rWgVP0BEgT91"
      },
      "source": [
        "## Gradient Bandits"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "twL9ltaQlgj-"
      },
      "source": [
        "In this task, you have to write a gradient bandits agent.\n",
        "\n",
        "The agent has these parameters:\n",
        "- `n_arms`: a number of arms, positive integer\n",
        "- `alpha`: a constant that regulates learning speed, float in range from 0 to 1\n",
        "- `fixed_rng`: a parameter used to verify the correctness of your solution, boolean\n",
        "\n",
        "Your agent should support changes in every parameter listed above, except `fixed_rng`. If you need to generate random values, use `self.np_random` instead of `np.random`.\n",
        "\n",
        "A small reminder:\n",
        "- probablity of action is computed according to a softmax formula: $P(A_t = a) = \\frac{e^{H_t(a)}}{\\sum_{b=1}^n e^{H_t(b)}}=\\pi_t(a)$, where softmax is computed as $\\sigma(x)_i=\\frac{e^{x_i}}{\\sum_{j=1}^{n}{e^{x_j}}}$\n",
        "- average reward can be computed like this: $\\bar{R}_t = \\frac{1}{t}(R_t-\\bar{R}_{t-1})$\n",
        "- the update formulas look like this:\n",
        "   - $H(A_t) \\gets H(A_t) + \\alpha(R_t-\\bar{R}_t)(1-\\pi_t(A_t))$ for selected action;\n",
        "   - $H(a) \\gets H(a) - \\alpha(R_t-\\bar{R}_t)\\pi_t(a)$ for all other actions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z_-d-p_Cy1kR"
      },
      "source": [
        "First, write a sofmax function. You will need this function to implement your agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O7pkD3iGutSM"
      },
      "outputs": [],
      "source": [
        "def softmax(x: np.ndarray) -> np.ndarray:\n",
        "    return None # The softmax value of x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zNA-kQ52zfNd"
      },
      "source": [
        "Now, you can write an agent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XE8xKWU1llr4"
      },
      "outputs": [],
      "source": [
        "class GradientBanditsAgent:\n",
        "    def __init__(self, n_arms: int, alpha: float, fixed_rng: bool = True):\n",
        "        self.np_random = np_random_generator(fixed_rng)\n",
        "\n",
        "        # Initialize agent's attributes\n",
        "        self.n_arms = None\n",
        "        self.alpha = None\n",
        "        self.H = None # Preferences\n",
        "        self.reward_avg = None # Average reward\n",
        "        self.t = None # Current time step\n",
        "\n",
        "    def select_action(self):\n",
        "        # Write an action selection algorithm\n",
        "        probs = None # Compute softmax probabilities\n",
        "        return None # Random action, selected according to probs\n",
        "\n",
        "    def update(self, action: int, reward: float):\n",
        "        # Write an update algorithm\n",
        "        self.t += None # Increase time step\n",
        "        self.reward_avg += None # Update average reward (mean)\n",
        "\n",
        "        probs = None # Compute softmax probabilities\n",
        "\n",
        "        # Update preferences according to the formulas\n",
        "        # Note that H includes H[action], so adapt the formula\n",
        "        self.H -= None\n",
        "        self.H[action] += None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AHP7kpMMiO_b"
      },
      "outputs": [],
      "source": [
        "solution4()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-GaeAryriR6Z"
      },
      "outputs": [],
      "source": [
        "check4(softmax, GradientBanditsAgent)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDSDGo0nL2Ds"
      },
      "source": [
        "## Testing the Algorithms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BuEi60Dsv5HT"
      },
      "source": [
        "Your last task for this challenge is to test your agents. To do this you will create a function, that takes these parameters:\n",
        "- `env_params`: environment parameters, dictionary\n",
        "- `agent_cls`: a class of an agent, class\n",
        "- `agent_params`: agent parameters, dictionary\n",
        "- `test_envs`: a number of test environments, int\n",
        "\n",
        "Your function should support changes in every parameter listed above.\n",
        "\n",
        "Also, your function should return two values:\n",
        "- `optimal_action_rate`: the rate of selecting an optimal action on each time step. For example, if on time step 5 optimal action was selected 6 times across all test environments, and there are 10 test environments, then the `optimal_action_rate[4]` is 0.6. You can check if the action was optimal by looking at `\"is_action_optimal\"` key in `info` dictionary, which is returned by a `step` function of an environment.\n",
        "- `average_return`: average return across all test environments\n",
        "\n",
        "**NOTE: use values 0..test_envs-1 for environment seeds**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T4LwqoTE0Y5c"
      },
      "source": [
        "First function should measure the performance of an agent in a stationary environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rw2kcCHBv4sy"
      },
      "outputs": [],
      "source": [
        "def test_agent(env_params: dict, agent_cls, agent_params: dict, test_envs: int):\n",
        "    # Initialize values that you will return\n",
        "    optimal_action_rate = None\n",
        "    average_return = None\n",
        "\n",
        "    # Initialize the environment\n",
        "    env = None\n",
        "\n",
        "    # Write a test loop for test_envs envs\n",
        "    for test_case in None:\n",
        "        agent = None # Create an agent from agent_cls and agent_params\n",
        "        pass # Reset the environment with seed=test_case\n",
        "\n",
        "\n",
        "        # Initialize values for results aggregation\n",
        "        t = None # Time step\n",
        "        total_reward = None\n",
        "        done = None\n",
        "\n",
        "        # Write a training loop\n",
        "        while None:\n",
        "            action = None # Select action according to the agent\n",
        "            _, reward, terminated, truncated, info = None # Do one step\n",
        "            pass # Update the agent's state\n",
        "            total_reward += None # Update total_reward\n",
        "\n",
        "            # Update optimal action stats\n",
        "            if None:\n",
        "                optimal_action_rate[t] += None\n",
        "            t += None # Increase time step\n",
        "\n",
        "            done = None # Update done status\n",
        "        average_return += None # Update average return\n",
        "\n",
        "    return optimal_action_rate, average_return"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "solution5()"
      ],
      "metadata": {
        "id": "RQG3yD6MpLco"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check5(test_agent)"
      ],
      "metadata": {
        "id": "Qst3_xxopMmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NLgl9B9W_xWY"
      },
      "source": [
        "Congratulations! This was your last task for this challenge. Now you can look at results of your work. Just run the cells below to see, how the algorithms perform. Don't worry about long execution times, this is OK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Dmg5tn1JvlS"
      },
      "source": [
        "### Stationary Environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NNL40IS5B1NJ"
      },
      "outputs": [],
      "source": [
        "env_params = {\n",
        "    \"id\": \"codefinityrl:MultiArmedBanditStationary-v0\",\n",
        "    \"max_episode_steps\": 1000,\n",
        "    \"n_arms\": 10,\n",
        "}\n",
        "\n",
        "eg_agent_params = {\n",
        "    \"n_arms\": 10,\n",
        "    \"epsilon\": 0,\n",
        "    \"alpha\": 0.1,\n",
        "    \"optimistic\": 1,\n",
        "    \"fixed_rng\": False,\n",
        "}\n",
        "\n",
        "ucb_agent_params = {\"n_arms\": 10, \"confidence\": 0.3}\n",
        "\n",
        "gb_agent_params = {\"n_arms\": 10, \"alpha\": 0.2, \"fixed_rng\": False}\n",
        "\n",
        "eg_optimal_action_rate, eg_average_return = test_agent(\n",
        "    env_params, EpsilonGreedyAgent, eg_agent_params, 1000\n",
        ")\n",
        "ucb_optimal_action_rate, ucb_average_return = test_agent(\n",
        "    env_params, UpperConfidenceBoundAgent, ucb_agent_params, 1000\n",
        ")\n",
        "gb_optimal_action_rate, gb_average_return = test_agent(\n",
        "    env_params, GradientBanditsAgent, gb_agent_params, 1000\n",
        ")\n",
        "\n",
        "fig, axs = plt.subplots(2, 1)\n",
        "fig.set_size_inches(9, 10)\n",
        "axs[0].set_title(\"Optimal action rates (Stationary env)\")\n",
        "axs[0].set_xlabel(\"Time step\")\n",
        "axs[0].set_ylabel(\"Rate\")\n",
        "axs[0].plot(eg_optimal_action_rate, color=\"r\", label=\"epsilon-greedy\")\n",
        "axs[0].plot(ucb_optimal_action_rate, color=\"g\", label=\"UCB\")\n",
        "axs[0].plot(gb_optimal_action_rate, color=\"b\", label=\"gradient bandits\")\n",
        "axs[0].legend(loc=\"best\")\n",
        "axs[1].set_title(\"Average return by algorithm\")\n",
        "axs[1].set_xlabel(\"Algorithm\")\n",
        "axs[1].set_ylabel(\"Average return\")\n",
        "bar = axs[1].bar(\n",
        "    [\"epsilon-greedy\", \"UCB\", \"gradient bandits\"],\n",
        "    [eg_average_return, ucb_average_return, gb_average_return],\n",
        "    color=[\"r\", \"g\", \"b\"],\n",
        ")\n",
        "axs[1].bar_label(bar, label_type=\"center\", fmt=\"%.2f\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TXAeS5tpJ2IP"
      },
      "source": [
        "### Dynamic Environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7RIUuzATIdak"
      },
      "outputs": [],
      "source": [
        "env_params = {\n",
        "    \"id\": \"codefinityrl:MultiArmedBanditDynamic-v0\",\n",
        "    \"max_episode_steps\": 5000,\n",
        "    \"n_arms\": 10,\n",
        "    \"drift_interval\": 500,\n",
        "}\n",
        "\n",
        "agent_params = {\"n_arms\": 10, \"epsilon\": 0.05, \"alpha\": 0.3, \"fixed_rng\": False}\n",
        "\n",
        "optimal_action_rate, average_return = test_agent(\n",
        "    env_params, EpsilonGreedyAgent, agent_params, 1000\n",
        ")\n",
        "\n",
        "plt.figure(figsize=(9, 5))\n",
        "plt.title(\"Optimal action rates (Dynamic env, epsilon-greedy)\")\n",
        "plt.xlabel(\"Time step\")\n",
        "plt.ylabel(\"Rate\")\n",
        "plt.plot(optimal_action_rate, color=\"r\")\n",
        "plt.show()\n",
        "print(f\"Average return: {average_return}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "toc_visible": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}