{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Challenge: Monte Carlo Methods"
      ],
      "metadata": {
        "id": "3A_YIgwVeLae"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Introduction"
      ],
      "metadata": {
        "id": "OQ6NwA5megOK"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this challenge, your goal is to implement two Monte Carlo control methods: on-policy and off-policy.\n",
        "\n",
        "You will:\n",
        "- Write a linear decay algorithm;\n",
        "- Write on-policy and off-policy Monte Carlo control algorithms;\n",
        "- Test your solutions in a custom environment;\n",
        "- Visually analyze action values and policy.\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.**"
      ],
      "metadata": {
        "id": "7ROwYGytehCo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Dependencies"
      ],
      "metadata": {
        "id": "GCWDirAEezC4"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Make sure to install and import all dependencies before you begin to work with the code"
      ],
      "metadata": {
        "id": "M22LA6CDe0w1"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JbUhGbwFw-1m"
      },
      "outputs": [],
      "source": [
        "!pip install numpy gymnasium git+https://github.com/Codefinity-ML/codefinity-rl.git"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import gymnasium as gym\n",
        "from codefinityrl.challenges.utils import VideoRecord, plot_values, plot_policy, np_random_generator\n",
        "from codefinityrl.challenges.mc import *"
      ],
      "metadata": {
        "id": "Cw_1OzlAxQdd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Environment"
      ],
      "metadata": {
        "id": "YM8YnrVcMQ7-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For this challenge, you will again use a custom environment **KeyAndChest**. Here is a small reminder of how it works.\n",
        "\n",
        "KeyAndChest environment represents a 2D maze with key and chest. To win, the agent has to find a key and use it to open a chest. Without a key, the chest can't be opened.\n",
        "\n",
        "The state space of this environment is represented by a tuple `(x, y, key)`, where\n",
        "- `x`: x coordinate, integer\n",
        "- `y`: y coordinate, integer\n",
        "- `key`: true if agent has key, otherwise - false, boolean\n",
        "\n",
        "The agent can take one of **four actions**:\n",
        "- UP - 0\n",
        "- DOWN - 1\n",
        "- LEFT - 2\n",
        "- RIGHT - 3\n",
        "\n",
        "**NOTE: this environment still has additional features, but you shouldn't use them in your solutions.**"
      ],
      "metadata": {
        "id": "a8-4EkflMS6o"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Linear Decay"
      ],
      "metadata": {
        "id": "6_RuqqMHD-SW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this task, your goal is to implement a simple linear decay algorithm. Later on, you will use it to decay the exploration rate $\\varepsilon$ in your on-policy and off-policy Monte Carlo control implementations.\n",
        "\n",
        "A linear decay starts from some big value and subtracts a fixed amount each time step, until it reaches a preset minimum.\n",
        "In code, these values are:\n",
        "- `steps`: number of steps, until value reaches minimum;\n",
        "- `high`: starting value;\n",
        "- `low`: minimum value.\n",
        "\n",
        "Therefore, the rate at which the value is decreased on each step is equal to $\\frac{high - low}{steps}$.\n",
        "\n",
        "Current value should decrease each time the `step` method is called, until the value reaches a preset minimum. After that, every call to step shouldn't change the value.\n",
        "\n",
        "`__call__` method should only return the current value."
      ],
      "metadata": {
        "id": "iLJNIvEmEJPp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Decay:\n",
        "    def __init__(self, steps: int, high: float = 1, low: float = 0):\n",
        "        # Initialize all object attributes\n",
        "        self.steps = None\n",
        "        self.high = None\n",
        "        self.low = None\n",
        "        self.decay_rate = None # Rate of decay\n",
        "        self.value = None # Current value\n",
        "\n",
        "    def step(self):\n",
        "        # Perform one decay step\n",
        "        self.value = None # Subtract decay_rate and clamp at low\n",
        "\n",
        "    def __call__(self):\n",
        "        # Return the current value when the instance is called\n",
        "        return self.value\n"
      ],
      "metadata": {
        "id": "H0IWCfOwOlWY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution1()"
      ],
      "metadata": {
        "id": "Jjf8V-nFDlhB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check1(Decay)"
      ],
      "metadata": {
        "id": "wDwcQmXoDn6M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can test, how your decay algorithm works by running the cell below."
      ],
      "metadata": {
        "id": "MgW62-K2DptP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "epsilon = Decay(10, 1, 0)\n",
        "for _ in range(15):\n",
        "    print(epsilon())\n",
        "    epsilon.step()"
      ],
      "metadata": {
        "id": "mOXU1ay6QDbJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## On-Policy Monte Carlo Control"
      ],
      "metadata": {
        "id": "Rji-4dJvJ05J"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this task, your goal is to implement an incremental on-policy Monte Carlo control algorithm.\n",
        "\n",
        "The pseudocode for incremental on-policy Monte Carlo control looked like this:\n",
        "<img src=\"https://codefinity-content-media-v2.s3.eu-west-1.amazonaws.com/courses/5f9425ad-7012-42d1-9fc0-d5998b1f5fd4/section-4/incremental-on-policy-monte-carlo-control.png\" />\n",
        "\n",
        "Notes on implementation:\n",
        "- Action values, visit counts, and policy for state $s$ stay uninitialized, until the state is first encountered. `init_state` method should be used to initialize the state;\n",
        "- Intial values for any state-action pair should be 0, except for terminal states, for which they are undefined;\n",
        "- Initial action for any non-terminal state should be 0. Terminal states should have the action undefined;\n",
        "- Generated episode should be a list of ($S_t$, $A_t$, $R_{t+1}$) tuples;\n",
        "- `update_policy` method should be used to update the policy for a state;\n",
        "- While two actions can have the same value, you should always select the action with lesser id when updating the policy;\n",
        "- Use `env.action_space.n` to get the number of available actions;\n",
        "- Use `self.np_random` if you need to generate random numbers."
      ],
      "metadata": {
        "id": "SuHCVAiYJ4gr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class OnPolicyMonteCarloControl:\n",
        "    def __init__(self, env: gym.Env):\n",
        "        self.np_random = np_random_generator(True)\n",
        "\n",
        "        # The environment is saved as the agent's attribute\n",
        "        self.env = env\n",
        "\n",
        "        # Initialize agent's policy, values and visit counts\n",
        "        self.policy = None\n",
        "        self.values = None\n",
        "        self.counts = None\n",
        "\n",
        "    def init_state(self, state):\n",
        "        # Initialize state if previously unseen\n",
        "        if None:\n",
        "            self.policy[state] = None\n",
        "            for action in None:\n",
        "                self.values[(state, action)] = None\n",
        "                self.counts[(state, action)] = None\n",
        "\n",
        "    def update_policy(self, state):\n",
        "        # Improve policy by choosing the action with the highest estimated value\n",
        "        best_action = None\n",
        "        best_action_value = None\n",
        "        for action in None:\n",
        "            if None:\n",
        "                best_action_value = None\n",
        "                best_action = None\n",
        "        self.policy[state] = None\n",
        "\n",
        "    def get_action(self, state, epsilon: float = 0):\n",
        "        # Choose an action using epsilon-greedy strategy\n",
        "        if None:\n",
        "            return None # Random action\n",
        "        else:\n",
        "            return None # Current optimal action\n",
        "\n",
        "    def get_episode(self, epsilon: float):\n",
        "        # Generate a single episode using current policy and epsilon\n",
        "\n",
        "        episode = None # Initialize an episode\n",
        "        state, _ = None # Get first state\n",
        "        done = None # Initialize termination check\n",
        "\n",
        "        # Add transitions until the episode is completed\n",
        "        while None:\n",
        "            pass # Initialize current state\n",
        "            action = None # Choose an action according to the policy\n",
        "            next_state, reward, terminated, truncated, _ = None # Perform a step\n",
        "\n",
        "            pass # Add transition to the episode\n",
        "            state = None # Move to a new state\n",
        "            done = None # Update termination check\n",
        "\n",
        "        return episode\n",
        "\n",
        "    def train(self, gamma: float, episodes: int, epsilon):\n",
        "        # Train the agent using on-policy Monte Carlo control\n",
        "\n",
        "        # Iterate over a given number of episodes\n",
        "        for _ in None:\n",
        "            episode = None # Generate an episode\n",
        "\n",
        "            G = None # Initialize return\n",
        "\n",
        "            # Go over every transition in reversed order\n",
        "            for state, action, reward in None:\n",
        "                G = None  # Update return\n",
        "\n",
        "                self.counts[(state, action)] += None # Increment visit count\n",
        "                self.values[(state, action)] += None # Update action value\n",
        "\n",
        "                pass # Update policy\n",
        "\n",
        "            pass # Decay epsilon\n",
        "\n",
        "    def __call__(self, state):\n",
        "        # Calling this agent with a state will return the action that the agent\n",
        "        # chooses according to the policy, or a random action if policy for\n",
        "        # the state is not defined\n",
        "        if state not in self.policy:\n",
        "            return self.np_random.integers(self.env.action_space.n)\n",
        "        return self.get_action(state)"
      ],
      "metadata": {
        "id": "rzpoOJbRGdqr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution2()"
      ],
      "metadata": {
        "id": "cZhdNL2YD9A4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check2(OnPolicyMonteCarloControl)"
      ],
      "metadata": {
        "id": "PJ6h0pHRD-Ma"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now run the cell below to train your agent."
      ],
      "metadata": {
        "id": "__emVhKpDhoE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\")\n",
        "on_policy_agent = OnPolicyMonteCarloControl(env)\n",
        "\n",
        "episodes = 10000\n",
        "epsilon = Decay(episodes, 1, 0.3)\n",
        "on_policy_agent.train(0.99, episodes, epsilon)"
      ],
      "metadata": {
        "id": "nEyOta6yTtot"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And run the cell below to verify that the agent solves the environment."
      ],
      "metadata": {
        "id": "x3a5JdaVEHxr"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\", render_mode=\"rgb_array\")\n",
        "state, _ = env.reset()\n",
        "\n",
        "video = VideoRecord(env.metadata[\"render_fps\"])\n",
        "video.add_frame(env.render())\n",
        "done = False\n",
        "while not done:\n",
        "    state, _, terminated, truncated, _ = env.step(on_policy_agent(state))\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "nwp3haDz_Ynh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to plot the value function and policy."
      ],
      "metadata": {
        "id": "FUfE6GdVENvA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(on_policy_agent.values)\n",
        "plot_policy(on_policy_agent.policy)"
      ],
      "metadata": {
        "id": "OPM8fsuzGc1D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As the action value heatmap shows, the agent clearly converges on the optimal path: state-action pairs along it have significantly higher values than the alternatives at the same state. Because the agent visits those optimal transitions far more often, its estimates there become much more accurate.\n",
        "\n",
        "You'll also notice that all action values are lower than one might initially expect. Since $\\varepsilon$-greedy exploration is baked into the Monte Carlo return estimates, every exploratory detour adds extra step penalties, pulling all values down.\n",
        "\n",
        "Finally, the greedy policy extracted from these action values is optimal in nearly every cell, but many choices remain \"close calls\". Off-path state-action pairs are visited so infrequently that their value estimates carry high variance, making them only marginally worse than the best action — and sometimes even better than the best action."
      ],
      "metadata": {
        "id": "nyUKwmouGGn9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Off-Policy Monte Carlo Control"
      ],
      "metadata": {
        "id": "cMuUWwEIPqcb"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this task, your goal is to implement an incremental off-policy Monte Carlo control algorithm.\n",
        "\n",
        "The pseudocode for incremental off-policy Monte Carlo control looked like this:\n",
        "<img src=\"https://codefinity-content-media-v2.s3.eu-west-1.amazonaws.com/courses/5f9425ad-7012-42d1-9fc0-d5998b1f5fd4/section-4/incremental-off-policy-monte-carlo-control.png\" />\n",
        "\n",
        "Notes on implementation:\n",
        "- Action values, sums of importance sampling ratios, and policy for state $s$ stay uninitialized, until the state is first encountered. `init_state` method should be used to initialize the state;\n",
        "- Intial values for any state-action pair should be 0, except for terminal states, for which they are undefined;\n",
        "- Initial action for any non-terminal state should be 0. Terminal states should have the action undefined;\n",
        "- Generated episode should be a list of ($S_t$, $A_t$, $R_{t+1}$) tuples;\n",
        "- `update_policy` method should be used to update the policy for a state;\n",
        "- While two actions can have the same value, you should always select the action with lesser id when updating the policy;\n",
        "- Use `env.action_space.n` to get the number of available actions;\n",
        "- Use `self.np_random` if you need to generate random numbers;\n",
        "- To obtain the probabilities required to compute the $\\frac{\\pi(a | s)}{b(a | s)}$ refer to the policy update rules in pseudocode;\n",
        "- Exit the processing of an episode if $W = 0$."
      ],
      "metadata": {
        "id": "9jfUZmSWPrvd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class OffPolicyMonteCarloControl:\n",
        "    def __init__(self, env: gym.Env):\n",
        "        self.np_random = np_random_generator(True)\n",
        "\n",
        "        # The environment is saved as the agent's attribute\n",
        "        self.env = env\n",
        "\n",
        "        # Initialize agent's policy, values and importance sampling ratio sums\n",
        "        self.policy = None\n",
        "        self.values = None\n",
        "        self.importances = None\n",
        "\n",
        "    def init_state(self, state):\n",
        "        # Initialize state if previously unseen\n",
        "        if None:\n",
        "            self.policy[state] = None\n",
        "            for action in None:\n",
        "                self.values[(state, action)] = None\n",
        "                self.importances[(state, action)] = None\n",
        "\n",
        "    def update_policy(self, state):\n",
        "        # Improve policy by choosing the action with the highest estimated value\n",
        "        best_action = None\n",
        "        best_action_value = None\n",
        "        for action in None:\n",
        "            if None:\n",
        "                best_action_value = None\n",
        "                best_action = None\n",
        "        self.policy[state] = None\n",
        "\n",
        "    def get_target_action(self, state):\n",
        "        # Choose an action greedily\n",
        "        return None\n",
        "\n",
        "    def get_behavior_action(self, state, epsilon: float):\n",
        "        # Choose an action using epsilon-greedy strategy\n",
        "        if None:\n",
        "            return None # Random action\n",
        "        else:\n",
        "            return None # Current optimal action\n",
        "\n",
        "\n",
        "    def get_episode(self, epsilon: float):\n",
        "        # Generate a single episode using current policy and epsilon\n",
        "\n",
        "        episode = None # Initialize an episode\n",
        "        state, _ = None # Get first state\n",
        "        done = None # Initialize termination check\n",
        "\n",
        "        # Add transitions until the episode is completed\n",
        "        while None:\n",
        "            pass # Initialize current state\n",
        "            action = None # Choose an action according to the policy\n",
        "            next_state, reward, terminated, truncated, _ = None # Perform a step\n",
        "\n",
        "            pass # Add transition to the episode\n",
        "            state = None # Move to a new state\n",
        "            done = None # Update termination check\n",
        "\n",
        "        return episode\n",
        "\n",
        "    def train(self, gamma: float, episodes: int, epsilon, target_epsilon: float):\n",
        "        # Train the agent using off-policy Monte Carlo control\n",
        "\n",
        "        # Iterate over a given number of episodes\n",
        "        for _ in None:\n",
        "            episode = None # Generate an episode\n",
        "\n",
        "            # Initialize return and importance sampling ratio\n",
        "            G = None\n",
        "            W = None\n",
        "\n",
        "            # Go over every transition in reversed order\n",
        "            for state, action, reward in None:\n",
        "                G = None # Update return\n",
        "                self.importances[(state, action)] += None # Update importance sum\n",
        "                self.values[(state, action)] = None # Update action value\n",
        "\n",
        "                pass # Update policy\n",
        "\n",
        "                # Update importance sampling ratio\n",
        "                # There are two general cases: either selected action\n",
        "                # is optimal or it isn't optimal\n",
        "                if None:\n",
        "                    W *= None # Update rule if action is optimal\n",
        "                else:\n",
        "                    W *= None # Update rule if action is not optimal\n",
        "\n",
        "                # Condition that exits if the importance sampling ratio\n",
        "                # reached zero to prevent numerical errors\n",
        "                if W == 0:\n",
        "                    break\n",
        "\n",
        "            epsilon.step()\n",
        "\n",
        "    def __call__(self, state):\n",
        "        # Calling this agent with a state will return the action that the agent\n",
        "        # chooses according to the policy, or a random action if policy for\n",
        "        # the state is not defined\n",
        "        if state not in self.policy:\n",
        "            return self.np_random.integers(self.env.action_space.n)\n",
        "        return self.get_action(state)"
      ],
      "metadata": {
        "id": "G6XRCVOc4QuY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "solution3()"
      ],
      "metadata": {
        "id": "JBvaspsKEarD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "check3(OffPolicyMonteCarloControl)"
      ],
      "metadata": {
        "id": "R2NZxrbCEcIB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now run the cell below to train your agent."
      ],
      "metadata": {
        "id": "F5wwLLCxEZzD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\")\n",
        "off_policy_agent = OffPolicyMonteCarloControl(env)\n",
        "\n",
        "episodes = 10000\n",
        "epsilon = Decay(episodes, 1, 0.3)\n",
        "off_policy_agent.train(0.99, episodes, epsilon, 0.01)"
      ],
      "metadata": {
        "id": "njQ2zET08zwr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "And run the cell below to verify that the agent solves the environment."
      ],
      "metadata": {
        "id": "nYJIvyvVEnFj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = gym.make(\"codefinityrl:KeyAndChest-v0\", render_mode=\"rgb_array\")\n",
        "state, _ = env.reset()\n",
        "\n",
        "video = VideoRecord(env.metadata[\"render_fps\"])\n",
        "video.add_frame(env.render())\n",
        "done = False\n",
        "while not done:\n",
        "    state, _, terminated, truncated, _ = env.step(off_policy_agent(state))\n",
        "    video.add_frame(env.render())\n",
        "    done = terminated or truncated\n",
        "video.play()"
      ],
      "metadata": {
        "id": "iMvr4b63872D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Run the cell below to plot the value function and policy."
      ],
      "metadata": {
        "id": "CoJs0xRRErvy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(off_policy_agent.values)\n",
        "plot_policy(off_policy_agent.policy)"
      ],
      "metadata": {
        "id": "o4-Df7vi9A5t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "When you compare this off-policy action value heatmap to the on-policy version, three things stand out:\n",
        "\n",
        "First of all, tighter gaps between action values. There's no single \"bright ribbon\" of clearly dominant actions, even though episodes are still sampled mostly along the known optimal paths. This is a natural result of using importance sampling.\n",
        "\n",
        "Second, the estimates are far more accurate. States near the goal or along the optimal return path now have action values that almost exactly match what you'd get from a dynamic programming solution. This unbiased precision is a direct consequence of weighted importance sampling.\n",
        "\n",
        "Lastly, as with the on-policy method, extracting the single best action at each state yields an optimal policy in nearly every cell. Only a handful of very rarely visited states still show suboptimal picks."
      ],
      "metadata": {
        "id": "DXO-VGNQR3MX"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Congratulations! You've completed the challenge."
      ],
      "metadata": {
        "id": "lS6BEY1hZrxQ"
      }
    }
  ]
}