week01_intro/crossentropy_method.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7XGyc-FCG35I"
      },
      "source": [
        "# Crossentropy method\n",
        "\n",
        "This notebook will teach you to solve reinforcement learning problems with crossentropy method. We'll follow-up by scaling everything up and using neural network policy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2jwz8moTG35K"
      },
      "outputs": [],
      "source": [
        "import sys, os\n",
        "if 'google.colab' in sys.modules and not os.path.exists('.setup_complete'):\n",
        "    !wget -q https://raw.githubusercontent.com/yandexdataschool/Practical_RL/master/setup_colab.sh -O- | bash\n",
        "    !touch .setup_complete\n",
        "\n",
        "# This code creates a virtual display to draw game images on.\n",
        "# It will have no effect if your machine has a monitor.\n",
        "if type(os.environ.get(\"DISPLAY\")) is not str or len(os.environ.get(\"DISPLAY\")) == 0:\n",
        "    !bash ../xvfb start\n",
        "    os.environ['DISPLAY'] = ':1'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oe7EKolvLC67"
      },
      "outputs": [],
      "source": [
        "!pip install gymnasium[toy_text]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ltjzx5AFG35K"
      },
      "outputs": [],
      "source": [
        "import gymnasium as gym\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "env = gym.make(\"Taxi-v3\", render_mode=\"rgb_array\")\n",
        "print(env.reset(seed=0))\n",
        "plt.imshow(env.render())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "O6waF01eG35L",
        "outputId": "8ca46444-7a7f-4091-c50c-0cab96a995e9"
      },
      "outputs": [],
      "source": [
        "n_states = env.observation_space.n\n",
        "n_actions = env.action_space.n\n",
        "\n",
        "print(f\"n_states={n_states}, n_actions={n_actions}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IHXnU2QWG35L"
      },
      "source": [
        "# Create stochastic policy\n",
        "\n",
        "This time our policy should be a probability distribution.\n",
        "\n",
        "```policy[s,a] = P(take action a | in state s)```\n",
        "\n",
        "Since we still use integer state and action representations, you can use a 2-dimensional array to represent the policy.\n",
        "\n",
        "Please initialize the policy __uniformly__, that is, probabililities of all actions should be equal."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "9qL5eW-rG35L"
      },
      "outputs": [],
      "source": [
        "def initialize_policy(n_states, n_actions):\n",
        "    <YOUR CODE: create an array to store action probabilities>\n",
        "\n",
        "    return policy\n",
        "\n",
        "\n",
        "policy = initialize_policy(n_states, n_actions)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "G1SeRRGgG35L"
      },
      "outputs": [],
      "source": [
        "assert type(policy) in (np.ndarray, np.matrix)\n",
        "assert np.allclose(policy, 1.0 / n_actions)\n",
        "assert np.allclose(np.sum(policy, axis=1), 1)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9zR1fCUrG35L"
      },
      "source": [
        "# Play the game\n",
        "\n",
        "Just like before, but we also record all states and actions we took."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "d4v8WmvlG35L"
      },
      "outputs": [],
      "source": [
        "def generate_session(env, policy, t_max=10**4):\n",
        "    \"\"\"\n",
        "    Play game until end or for t_max ticks.\n",
        "    :param policy: an array of shape [n_states,n_actions] with action probabilities\n",
        "    :returns: list of states, list of actions and sum of rewards\n",
        "    \"\"\"\n",
        "    states, actions = [], []\n",
        "    total_reward = 0.0\n",
        "\n",
        "    s, _ = env.reset()\n",
        "\n",
        "    for t in range(t_max):\n",
        "        # Hint: you can use np.random.choice for sampling action\n",
        "        # https://numpy.org/doc/stable/reference/random/generated/numpy.random.choice.html\n",
        "\n",
        "        a = <YOUR CODE: sample action from policy>\n",
        "\n",
        "        new_s, r, terminated, truncated, _ = env.step(a)\n",
        "\n",
        "        # Record information we just got from the environment.\n",
        "        states.append(s)\n",
        "        actions.append(a)\n",
        "        total_reward += r\n",
        "\n",
        "        s = new_s\n",
        "        if terminated or truncated:\n",
        "            break\n",
        "\n",
        "    return states, actions, total_reward\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "a1EUUZ29G35M"
      },
      "outputs": [],
      "source": [
        "s, a, r = generate_session(env, policy)\n",
        "assert type(s) == type(a) == list\n",
        "assert len(s) == len(a)\n",
        "assert type(r) in [float, np.float64]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_5YEDTKnG35M"
      },
      "outputs": [],
      "source": [
        "# let's see the initial reward distribution\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "\n",
        "sample_rewards = [generate_session(env, policy, t_max=1000)[-1] for _ in range(200)]\n",
        "\n",
        "plt.hist(sample_rewards, bins=20)\n",
        "plt.vlines([np.percentile(sample_rewards, 50)], [0], [100], label=\"50'th percentile\", color='green')\n",
        "plt.vlines([np.percentile(sample_rewards, 90)], [0], [100], label=\"90'th percentile\", color='red')\n",
        "plt.legend()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EeWtL3F5G35M"
      },
      "source": [
        "### Crossentropy method steps"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "00WWzr0KG35N"
      },
      "outputs": [],
      "source": [
        "def select_elites(states_batch, actions_batch, rewards_batch, percentile):\n",
        "    \"\"\"\n",
        "    Select states and actions from games that have rewards >= percentile\n",
        "    :param states_batch: list of lists of states, states_batch[session_i][t]\n",
        "    :param actions_batch: list of lists of actions, actions_batch[session_i][t]\n",
        "    :param rewards_batch: list of rewards, rewards_batch[session_i]\n",
        "\n",
        "    :returns: elite_states,elite_actions, both 1D lists of states and respective actions from elite sessions\n",
        "\n",
        "    Please return elite states and actions in their original order\n",
        "    [i.e. sorted by session number and timestep within session]\n",
        "\n",
        "    If you are confused, see examples below. Please don't assume that states are integers\n",
        "    (they will become different later).\n",
        "    \"\"\"\n",
        "\n",
        "    reward_threshold = <YOUR CODE: compute minimum reward for elite sessions. Hint: use np.percentile()>\n",
        "\n",
        "    elite_states = <YOUR CODE>\n",
        "    elite_actions = <YOUR CODE>\n",
        "\n",
        "    return elite_states, elite_actions\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "teOLBGojG35N"
      },
      "outputs": [],
      "source": [
        "states_batch = [\n",
        "    [1, 2, 3],     # game1\n",
        "    [4, 2, 0, 2],  # game2\n",
        "    [3, 1],        # game3\n",
        "]\n",
        "\n",
        "actions_batch = [\n",
        "    [0, 2, 4],     # game1\n",
        "    [3, 2, 0, 1],  # game2\n",
        "    [3, 3],        # game3\n",
        "]\n",
        "rewards_batch = [\n",
        "    3,  # game1\n",
        "    4,  # game2\n",
        "    5,  # game3\n",
        "]\n",
        "\n",
        "test_result_0 = select_elites(states_batch, actions_batch, rewards_batch, percentile=0)\n",
        "test_result_30 = select_elites(\n",
        "    states_batch, actions_batch, rewards_batch, percentile=30\n",
        ")\n",
        "test_result_90 = select_elites(\n",
        "    states_batch, actions_batch, rewards_batch, percentile=90\n",
        ")\n",
        "test_result_100 = select_elites(\n",
        "    states_batch, actions_batch, rewards_batch, percentile=100\n",
        ")\n",
        "\n",
        "assert np.all(test_result_0[0] == [1, 2, 3, 4, 2, 0, 2, 3, 1]) and np.all(\n",
        "    test_result_0[1] == [0, 2, 4, 3, 2, 0, 1, 3, 3]\n",
        "), \"For percentile 0 you should return all states and actions in chronological order\"\n",
        "assert np.all(test_result_30[0] == [4, 2, 0, 2, 3, 1]) and np.all(\n",
        "    test_result_30[1] == [3, 2, 0, 1, 3, 3]\n",
        "), \"For percentile 30 you should only select states/actions from two first\"\n",
        "assert np.all(test_result_90[0] == [3, 1]) and np.all(\n",
        "    test_result_90[1] == [3, 3]\n",
        "), \"For percentile 90 you should only select states/actions from one game\"\n",
        "assert np.all(test_result_100[0] == [3, 1]) and np.all(\n",
        "    test_result_100[1] == [3, 3]\n",
        "), \"Please make sure you use >=, not >. Also double-check how you compute percentile.\"\n",
        "\n",
        "print(\"Ok!\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wYLBHFFwG35N"
      },
      "outputs": [],
      "source": [
        "def get_new_policy(elite_states, elite_actions):\n",
        "    \"\"\"\n",
        "    Given a list of elite states/actions from select_elites,\n",
        "    return a new policy where each action probability is proportional to\n",
        "\n",
        "        policy[s_i,a_i] ~ #[occurrences of s_i and a_i in elite states/actions]\n",
        "\n",
        "    Don't forget to normalize the policy to get valid probabilities and handle the 0/0 case.\n",
        "    For states that you never visited, use a uniform distribution (1/n_actions for all states).\n",
        "\n",
        "    :param elite_states: 1D list of states from elite sessions\n",
        "    :param elite_actions: 1D list of actions from elite sessions\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    new_policy = np.zeros([n_states, n_actions])\n",
        "\n",
        "    <YOUR CODE: set probabilities for actions given elite states & actions>\n",
        "    # Don't forget to set 1/n_actions for all actions in unvisited states.\n",
        "\n",
        "    return new_policy\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I1VbNcpoG35O"
      },
      "outputs": [],
      "source": [
        "elite_states = [1, 2, 3, 4, 2, 0, 2, 3, 1]\n",
        "elite_actions = [0, 2, 4, 3, 2, 0, 1, 3, 3]\n",
        "\n",
        "new_policy = get_new_policy(elite_states, elite_actions)\n",
        "\n",
        "assert np.isfinite(\n",
        "    new_policy\n",
        ").all(), \"Your new policy contains NaNs or +-inf. Make sure you don't divide by zero.\"\n",
        "assert np.all(\n",
        "    new_policy >= 0\n",
        "), \"Your new policy can't have negative action probabilities\"\n",
        "assert np.allclose(\n",
        "    new_policy.sum(axis=-1), 1\n",
        "), \"Your new policy should be a valid probability distribution over actions\"\n",
        "\n",
        "reference_answer = np.array(\n",
        "    [\n",
        "        [1.0, 0.0, 0.0, 0.0, 0.0],\n",
        "        [0.5, 0.0, 0.0, 0.5, 0.0],\n",
        "        [0.0, 0.33333333, 0.66666667, 0.0, 0.0],\n",
        "        [0.0, 0.0, 0.0, 0.5, 0.5],\n",
        "    ]\n",
        ")\n",
        "assert np.allclose(new_policy[:4, :5], reference_answer)\n",
        "\n",
        "print(\"Ok!\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WVvZq1fSG35O"
      },
      "source": [
        "# Training loop\n",
        "Generate sessions, select N best and fit to those."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3CmH7Aj4G35O"
      },
      "outputs": [],
      "source": [
        "from IPython.display import clear_output\n",
        "\n",
        "\n",
        "def show_progress(rewards_batch, log, percentile, reward_range=[-990, +10]):\n",
        "    \"\"\"\n",
        "    A convenience function that displays training progress.\n",
        "    No cool math here, just charts.\n",
        "    \"\"\"\n",
        "\n",
        "    mean_reward = np.mean(rewards_batch)\n",
        "    threshold = np.percentile(rewards_batch, percentile)\n",
        "    log.append([mean_reward, threshold])\n",
        "\n",
        "    plt.figure(figsize=[8, 4])\n",
        "    plt.subplot(1, 2, 1)\n",
        "    plt.plot(list(zip(*log))[0], label=\"Mean rewards\")\n",
        "    plt.plot(list(zip(*log))[1], label=\"Reward thresholds\")\n",
        "    plt.legend()\n",
        "    plt.grid()\n",
        "\n",
        "    plt.subplot(1, 2, 2)\n",
        "    plt.hist(rewards_batch, range=reward_range)\n",
        "    plt.vlines(\n",
        "        [np.percentile(rewards_batch, percentile)],\n",
        "        [0],\n",
        "        [100],\n",
        "        label=\"percentile\",\n",
        "        color=\"red\",\n",
        "    )\n",
        "    plt.legend()\n",
        "    plt.grid()\n",
        "    clear_output(True)\n",
        "    print(\"mean reward = %.3f, threshold=%.3f\" % (mean_reward, threshold))\n",
        "    plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tz0Yd964G35O"
      },
      "outputs": [],
      "source": [
        "# reset policy just in case\n",
        "policy = initialize_policy(n_states, n_actions)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-cNq5hndG35O"
      },
      "outputs": [],
      "source": [
        "n_sessions = 250  # sample this many sessions\n",
        "percentile = 50  # discard this percentage of sessions with lowest rewards\n",
        "learning_rate = 0.5  # how quickly the policy is updated, on a scale from 0 to 1\n",
        "\n",
        "log = []\n",
        "\n",
        "for i in range(100):\n",
        "    %time sessions = [ <YOUR CODE: generate a list of n_sessions new sessions> ]\n",
        "\n",
        "    states_batch, actions_batch, rewards_batch = zip(*sessions)\n",
        "\n",
        "    elite_states, elite_actions = <YOUR CODE: select elite states & actions>\n",
        "\n",
        "    new_policy = <YOUR CODE: compute new policy>\n",
        "\n",
        "    policy = learning_rate * new_policy + (1 - learning_rate) * policy\n",
        "\n",
        "    # display results on chart\n",
        "    show_progress(rewards_batch, log, percentile)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K5LIoVTuG35O"
      },
      "source": [
        "### Reflecting on results\n",
        "\n",
        "You may have noticed that the taxi problem quickly converges from less than -1000 to a near-optimal score and then descends back into -50/-100. This is in part because the environment has some innate randomness. Namely, the starting points of passenger/driver change from episode to episode.\n",
        "\n",
        "In case CEM failed to learn how to win from one distinct starting point, it will simply discard it because no sessions from that starting point will make it into the \"elites\".\n",
        "\n",
        "To mitigate that problem, you can either reduce the threshold for elite sessions (duct tape way) or change the way you evaluate strategy (theoretically correct way). For each starting state, you can sample an action randomly, and then evaluate this action by running _several_ games starting from it and averaging the total reward. Choosing elite sessions with this kind of sampling (where each session's reward is counted as the average of the rewards of all sessions with the same starting state and action) should improve the performance of your policy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ubIHBgQ-G35O"
      },
      "source": [
        "\n",
        "### You're not done yet!\n",
        "\n",
        "Go to [`./deep_crossentropy_method.ipynb`](./deep_crossentropy_method.ipynb) for a more serious task"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}