Primal_Dual_Students.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uWN_96PKsM_j"
   },
   "source": [
    "Author: Germano Gabbianelli\n",
    "\n",
    "License: GNU GPL v3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Install and import the required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install numpy scipy matplotlib pygame gymnasium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "n3L_uOwcsM_k",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import gymnasium as gym\n",
    "from gymnasium.envs.toy_text.frozen_lake import generate_random_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-9PSrHLXsM_q"
   },
   "source": [
    "## Configure the Random number Generators\n",
    "\n",
    "We set an initial seed to ensure reproducibility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tnR0CgT6o7v0",
    "tags": []
   },
   "outputs": [],
   "source": [
    "RANDOM_SEED = 666"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use a seed sequence to initialize all the possible random number generators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "seed_seq = np.random.SeedSequence(RANDOM_SEED)\n",
    "\n",
    "MAP_SEED, ENV_SEED, ALG_SEED = seed_seq.spawn(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unfortunately, for some reason, gym checks that the seeds are `int`s,\n",
    "so we need this small trick"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "MAP_SEED = int(MAP_SEED.generate_state(1)[0])\n",
    "ENV_SEED = int(ENV_SEED.generate_state(1)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Configure the FrozenLake Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can generate a random map using `generate_random_map`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "R-dR83cgsM_q",
    "outputId": "7eb93f3d-0d1f-4d22-c44d-e3d2eeea9b79",
    "tags": []
   },
   "outputs": [],
   "source": [
    "M = MAP_SIZE = 5\n",
    "FROZEN_PROBABILITY=0.9\n",
    "\n",
    "map = generate_random_map(\n",
    "    size=MAP_SIZE,\n",
    "    p=FROZEN_PROBABILITY,\n",
    "    seed=ENV_SEED,\n",
    ")\n",
    "print(map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or we can specify one ourselves:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# rename this variable to `map` to use this map\n",
    "# everything works (should) the same but you'll have\n",
    "# a bit longer training times\n",
    "_map = [\n",
    "    \"FFFFF\",\n",
    "    \"FHFFF\",\n",
    "    \"FHFHF\",\n",
    "    \"FHHHF\",\n",
    "    \"FFSHG\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us initialize the gym environment and print it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "env = gym.make(\n",
    "    \"FrozenLake-v1\",\n",
    "    is_slippery=False,\n",
    "    render_mode=\"rgb_array\",\n",
    "    desc=map,\n",
    ")\n",
    "\n",
    "\n",
    "env.reset(seed=ENV_SEED) # set the random seed\n",
    "\n",
    "env_image = env.render()\n",
    "plt.imshow(env_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define an Interface for our Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us define a base class for all our algorithms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import abc\n",
    "\n",
    "class Algorithm(abc.ABC):\n",
    "\n",
    "    def __init__(self, env, df=0.95, seed=None):\n",
    "        self.env = env\n",
    "        self.df = df\n",
    "        \n",
    "        # useful as a shortcut\n",
    "        self.S = env.observation_space.n\n",
    "        self.A = env.action_space.n\n",
    "\n",
    "        # initialize the random number generator\n",
    "        self.rng = np.random.default_rng(seed)\n",
    "\n",
    "\n",
    "    def step(self, t):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @abc.abstractproperty\n",
    "    def policy(self):\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also define some utility functions/classes (hidden by default)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qLPzbBPq7KBn",
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Utility functions/classes\n",
    "\n",
    "def plot_policy(policy):\n",
    "\n",
    "    U = (policy[:, 2] - policy[:, 0]).reshape(M, M)\n",
    "    V = (policy[:, 3] - policy[:, 1]).reshape(M, M)\n",
    "\n",
    "    x = np.arange(M) + 0.5\n",
    "    X, Y = np.meshgrid(x, x[::-1])\n",
    "\n",
    "    #U = np.ones(25).reshape(5,5)\n",
    "    #V = U\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "\n",
    "    ax.imshow(env_image, extent=[0, M, 0, M])\n",
    "\n",
    "    ax.quiver(X, Y, U,V, angles='xy', scale_units='xy', scale=2)\n",
    "    ax.set(xlim=[0,M], ylim=[0,M], aspect=\"equal\")\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "    \n",
    "def compute_P(env):\n",
    "    S = env.observation_space.n\n",
    "    A = env.action_space.n\n",
    "\n",
    "    P = P = np.zeros((S, A, S))\n",
    "\n",
    "    for (state, state_data) in env.P.items():\n",
    "        for (action, next_data) in state_data.items():\n",
    "            for (prob, next_state, reward, terminated) in next_data:\n",
    "                P[state, action, next_state] = prob\n",
    "\n",
    "    return P\n",
    "\n",
    "\n",
    "def compute_R(env):\n",
    "    S = env.observation_space.n\n",
    "    A = env.action_space.n\n",
    "\n",
    "    R = np.zeros((S, A, S))\n",
    "\n",
    "    for (state, state_data) in env.P.items():\n",
    "        for (action, next_data) in state_data.items():\n",
    "            for (prob, next_state, reward, terminated) in next_data:\n",
    "                R[state, action, next_state] = reward\n",
    "\n",
    "    return R\n",
    "\n",
    "\n",
    "class InvalidatePolicy:\n",
    "    def __set_name__(self, owner, name):\n",
    "        self.name = name\n",
    "\n",
    "    def __set__(self, instance, value):\n",
    "        instance.__dict__[self.name] = value\n",
    "        instance.__dict__.pop(\"policy\", None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QROz1W64Al9R"
   },
   "source": [
    "## Linear Programming Formulation Recap\n",
    "\n",
    "The Reinforcement Learning Problem can be formulated with the following two linear programs:\n",
    "\n",
    "$\\newcommand{\\One}[1]{1\\!\\!1\\{#1\\}}$\n",
    "$$\n",
    "    \\begin{aligned}[c]\n",
    "        & \\max_{\\mu\\succeq 0}\\; \\langle \\mu, r \\rangle \\\\\n",
    "        & \\textrm{subject to}\\ E^\\intercal \\mu =(1-\\gamma)\\nu_0 + \\gamma P^\\intercal \\mu \\\\\n",
    "    \\end{aligned}\\quad\\;\\;\n",
    "    \\begin{aligned}[c]\n",
    "    &\\min_{v}\\; (1-\\gamma) \\langle \\nu_0, v\\rangle \\\\\n",
    "    &\\textrm{subject to}\\ Ev \\geq r+\\gamma Pv,\n",
    "    \\end{aligned}\n",
    "$$\n",
    "\n",
    "From which we can derive the Lagrangian\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathcal{L}(v, \\mu) &= (1-\\gamma)\\langle\\nu_0, v\\rangle + \\langle\\mu, r+\\gamma Pv - Ev\\rangle \\\\\n",
    " &= \\langle\\mu, r \\rangle + \\langle v, (1-\\gamma)\\nu_0 + \\gamma P^\\intercal \\mu - E^\\intercal \\mu \\rangle .\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "We then use the gradients of the Lagrangian as losses for two alternativing mirror descent / ascent instances.\n",
    "\n",
    "<!--\n",
    "Now consider the random variable $W=(X_0, X, A, R, X')$, where $X_0\\sim \\nu_0$,\n",
    "$(X,A)\\sim\\mu$, $R=r(X,A)$, and $X'\\sim p(\\cdot\\mid X,A)$. Then, we can define\n",
    "the unbiased estimators\n",
    "\n",
    "*   $\\hat{\\nu}_0(x) = \\One{X^0=x}$,\n",
    "*   $\\hat{r}(x,a)= \\One{X=x, A=a}R$\n",
    "*   $\\hat{p}(x'\\mid x,a)=\\One{X'=x'\\mid X=x, A=a}$ -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Exercise: compute the gradients\n",
    "\n",
    "Compute the gradients of the Lagrangian wrt $v$ and $\\mu$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## A base class for our Tabular algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now going to write a base class for all our tabular algorithms.\n",
    "It initializes the variables used in our algorithms (`v` and `mu`),\n",
    "and their learning rate. It also defines a `mu_sum` variable, because\n",
    "the output policy is extracted from the sum of the `mu`s (Do you remember\n",
    "this deatil from yesterday's lesson?).\n",
    "\n",
    "It also makes avaiable to subclasses the attributes `P`, `r` and `nu0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oKT_JB-LW9og",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from functools import cached_property\n",
    "\n",
    "\n",
    "class TabularBase(Algorithm):\n",
    "\n",
    "    # Makes sure the policy is deleted every time mu_sum is assigned\n",
    "    mu_sum = InvalidatePolicy()\n",
    "\n",
    "    def __init__(self, env, **kwds):\n",
    "        lr_v = lr_mu = kwds.pop(\"lr\", 0.01)\n",
    "        self.lr_v = kwds.pop(\"lr_v\", lr_v)\n",
    "        self.lr_mu = kwds.pop(\"lr_mu\", lr_mu)\n",
    "        \n",
    "        super().__init__(env, **kwds)\n",
    "\n",
    "        # initialize variables\n",
    "        self.v = np.zeros(self.S)\n",
    "        self.mu = np.ones((self.S, self.A)) / (self.S * self.A)\n",
    "        self.mu_sum = self.mu\n",
    "\n",
    "    @cached_property\n",
    "    # @cached_property computes the attribute on first access and then caches it\n",
    "    def policy(self):\n",
    "        ###################################\n",
    "        # Exercise                        #\n",
    "        ###################################\n",
    "        # What output policy should we return?\n",
    "        # Hint: extract it from self.mu_sum\n",
    "        return ...\n",
    "\n",
    "    @cached_property\n",
    "    def P(self):\n",
    "        return compute_P(self.env)\n",
    "\n",
    "    @cached_property\n",
    "    def r(self):\n",
    "        R = compute_R(self.env)\n",
    "        return (self.P * R).sum(2)\n",
    "\n",
    "    @property\n",
    "    def nu0(self):\n",
    "        return self.env.initial_state_distrib\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a small utility class to log the value of the policies produced by our\n",
    "algorithm during training (hidden by default):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6WHBQUIXapkV",
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class LoggingMixin:\n",
    "    \"Utility class which evaluates the current policy at each step\\\n",
    "     and saves the values in `_rewards`.\"\n",
    "    \n",
    "    def __init__(self, env, *args, log_every=10, **kwds):\n",
    "        super().__init__(env, *args, **kwds)\n",
    "\n",
    "        self._log_every = log_every\n",
    "        self._I = np.eye(self.S)\n",
    "        self._rewards = []\n",
    "\n",
    "    def evaluate_policy(self, policy):\n",
    "        P = (self.P * policy.reshape(self.S, self.A, 1)).sum(1)\n",
    "        r = (self.r * policy).sum(1)\n",
    "\n",
    "        v = np.linalg.inv(self._I - self.df * P) @ r\n",
    "        return (1-self.df) * self.nu0 @ v\n",
    "\n",
    "    def step(self, t):\n",
    "        super().step(t)\n",
    "\n",
    "        if t % self._log_every == 0:\n",
    "            r = self.evaluate_policy(self.policy)\n",
    "            self._rewards.append(r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute the exact solution with a solver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before proceeding to implement some algorithm ourselves, it is insightful to use a Linear Programming solver (i.e. `scipy.optimize.linprog`) to find a solution\n",
    "for our problem.\n",
    "\n",
    "The solver requires the objective to be in the form $\\langle x,c \\rangle$, where $x$ is the variable to *minimize*,\n",
    "and the constraints to be in the form $Ax = b$.\n",
    "\n",
    "We can therefore call the `linprog` function passing the appropriate value for the parameters `c`, `A_eq` and `b_eq`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy.optimize import linprog\n",
    "\n",
    "\n",
    "class PrimalSolver(TabularBase):\n",
    "    def __init__(self, *args, **kwds):\n",
    "        super().__init__(*args, **kwds)\n",
    "        \n",
    "        ## First, we compute the E matrix\n",
    "        E = np.zeros((self.S, self.A, self.S))\n",
    "        ii = np.arange(self.S)\n",
    "        \n",
    "        E[ii, :, ii] = 1\n",
    "        \n",
    "        self.E = E\n",
    "        \n",
    "    \n",
    "    def step(self, t=0):\n",
    "        \n",
    "        SA = self.S * self.A\n",
    "        \n",
    "        r = self.r.reshape(SA)\n",
    "        P = self.P.reshape(SA, self.S)\n",
    "        E = self.E.reshape(SA, self.S)\n",
    "        \n",
    "        self._primal = primal = linprog(\n",
    "            c = ..., # keep in mind linprog only supports minimization!\n",
    "            A_eq = ...,\n",
    "            b_eq = ...,\n",
    "            bounds = ...,\n",
    "        )\n",
    "        \n",
    "        mu = primal.x.reshape(self.S, self.A)\n",
    "        nu = mu.sum(1)\n",
    "        \n",
    "        # We set the policy to be uniform for states\n",
    "        # not visited by the deterministic optimal policy\n",
    "        # produced by the solver\n",
    "        mu[nu == 0] = 1 / self.A\n",
    "        \n",
    "        self.mu = self.mu_sum = mu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = PrimalSolver(env)\n",
    "agent.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(agent._primal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_policy(agent.policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Let's do the same with the Dual"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now try to solve the dual problem. The constraints must be expressed in the form $Ax \\leq b$, and the name of the parameters are now `A_ub` and `b_ub`.\n",
    "\n",
    "Question: how do we extract the optimal policy from $v$?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class DualSolver(TabularBase):\n",
    "    \n",
    "    def __init__(self, *args, **kwds):\n",
    "        super().__init__(*args, **kwds)\n",
    "        \n",
    "        E = np.zeros((self.S, self.A, self.S))\n",
    "        ii = np.arange(self.S)\n",
    "        \n",
    "        E[ii, :, ii] = 1\n",
    "        \n",
    "        self.E = E\n",
    "    \n",
    "    def step(self, t=0):\n",
    "        \n",
    "        SA = self.S * self.A\n",
    "        r = self.r.reshape(SA)\n",
    "        P = self.P.reshape(SA, self.S)\n",
    "        E = self.E.reshape(SA, self.S)\n",
    "\n",
    "\n",
    "        dual = linprog(\n",
    "            c=...,\n",
    "            A_ub=...,\n",
    "            b_ub=...,\n",
    "            bounds=...,\n",
    "        )\n",
    "        print(dual)\n",
    "        self._dual = dual       \n",
    "\n",
    "        self.v = dual.x.reshape(self.S)\n",
    "        self.q = (r + self.df * P @ self.v).reshape(self.S, self.A)\n",
    "        \n",
    "                \n",
    "    @cached_property\n",
    "    def policy(self):\n",
    "        pi = np.zeros((self.S, self.A))\n",
    "        ii = np.arange(self.S)\n",
    "        \n",
    "        pi[ii, self.q.argmax(1)] = 1\n",
    "        return pi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SA = agent.S * agent.A\n",
    "E = agent.E.reshape(SA, agent.S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = DualSolver(env)\n",
    "agent.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_policy(agent.policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Our first algorithm: Distributional Primal Dual"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now going to implement a primal-dual method which requires the use of a distributional model of the environment (i.e. access to `P`, `r` and `nu0`).\n",
    "\n",
    "We are going to need to remember the gradients of the Lagrangian with respect to $\\mu$ and $v$ and the update rules of gradient descent and exponential weights.\n",
    "\n",
    "**Note**: Assuming rewards bounded in $[0,1]$, we have that for any value function and any state $s$,\n",
    "$$\n",
    "0 \\leq v(s) \\leq (1-\\gamma)^{-1}\n",
    "$$\n",
    "\n",
    "Therefore, we will project the $v$ computed at each step into the set $[0, (1-\\gamma)^{-1}]^{|\\S|}$.\n",
    "\n",
    "\n",
    "**Hint**: If you need a little help with the definitons of gradient descent or exponential weights, expand the next cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "source": [
    "Given a learning rate $\\eta$. The updated of gradient descent is given by\n",
    "$$\n",
    "x \\gets x - \\eta \\nabla_x f\n",
    "$$\n",
    "\n",
    "Exponential Weights:\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\tilde{x} &\\gets x \\circ \\exp(\\eta \\nabla_x f)\\\\\n",
    "x &\\gets \\tilde{x} / (E^\\top \\tilde{x})\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "guxXlBykmH0O",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class DistribPD(TabularBase):\n",
    "\n",
    "    def grad_mu(self, v):\n",
    "        # Compute the gradient of the Lagrangian wrt to mu\n",
    "        # evaluated at point `v`\n",
    "        return ...\n",
    "\n",
    "    def grad_v(self, mu):\n",
    "        # compute the gradient of the Lagrangian wrt to v\n",
    "        # evaluated at point `mu`\n",
    "        return ...\n",
    "\n",
    "    def step_v(self):\n",
    "        ## Do a gradient descent update step on v\n",
    "        ## using learning rate self.lr_v\n",
    "        \n",
    "        # compute the gradient evaluated at the current point\n",
    "        grad = self.grad_v(self.mu)\n",
    "        \n",
    "        # do the gradient descent update (do not assign to self.v,\n",
    "        # as it is already done in the `step` method)\n",
    "        v = ...\n",
    "        \n",
    "        # remember to project v into the set [0, 1/(1-gamma)]^|S|\n",
    "        # ask for a hint if you are stuck\n",
    "        \n",
    "        return v\n",
    "\n",
    "\n",
    "    def step_mu(self):\n",
    "        ## Do an exponential weight update step on mu\n",
    "        ## using learning rate self.lr_mu\n",
    "        \n",
    "        # compute the gradient evaluated at the current point\n",
    "        grad = self.grad_mu(self.v)\n",
    "        \n",
    "        # implement the exponential weights update step\n",
    "        mu = ...\n",
    "        \n",
    "        return mu\n",
    "\n",
    "    def step(self, t):\n",
    "        v = agent.step_v()\n",
    "        mu = agent.step_mu()\n",
    "\n",
    "        self.v = v\n",
    "        self.mu = mu\n",
    "        self.mu_sum += mu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the `LoggingMixin` to make sure that the value of the policy produced by our algorithm is\n",
    "logged (i.e. saved to the `_rewards` attribute of the agent)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class LoggingDistribPD(LoggingMixin, DistribPD):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are finally ready to run our algorithm. Experiment with the number of learning steps and the learning rate, until the algorithm converges and achives the optimal value (Hint: look up the optimal value in the output of the solver above)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9vaPa0NpAutN",
    "tags": []
   },
   "outputs": [],
   "source": [
    "n = 20_000\n",
    "lr = 1/np.sqrt(n)\n",
    "\n",
    "agent = LoggingDistribPD(env, lr=lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sRmrwz1a6-0_",
    "tags": []
   },
   "outputs": [],
   "source": [
    "for t in range(n):\n",
    "    agent.step(t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the values of the policies produced at each step by the algorithm ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 452
    },
    "id": "4X9AAuEcki6h",
    "outputId": "84cb0c5b-29d0-4c26-8b47-c12440c6913f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.plot(agent._rewards)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and the output policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 440
    },
    "id": "8a9y_3OSCAIG",
    "outputId": "ae000512-6e41-47b6-feb7-fe439b2ba05d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_policy(agent.policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Second Algorithm: Tabular Primal Dual with Generative Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now relax our assumption of having access to a distributional model. Let's implement a primal-dual method only requiring access to a generative model (i.e. given any state-action pair $s,a$ it returns a sample $s'\\sim p(\\cdot | s,a)$).\n",
    "\n",
    "### Exercise: rewrite the gradients as an expectation\n",
    "\n",
    "1. Can you re-write the gradient of the Lagrangian wrt $v$ as an expectation over $\\mu$ ?\n",
    "2. Can you re-write the gradient of the Lagrangian wrt $\\mu$ as an expectation over the uniform policy?\n",
    "\n",
    "Expand the next cell to find the solution to 1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "source": [
    "If $X$ and $A$ are sampled according to $\\mu$. Then,\n",
    "\n",
    "$$\n",
    "\\newcommand{\\S}{\\mathcal{S}}\n",
    "\\newcommand{\\A}{\\mathcal{A}}\n",
    "\\nabla_v \\mathcal{L}(v, \\mu)_x = \\mathbb{E}[ \\One{X^0 = x} + \\gamma \\One{X' = x} - \\One{X = x} ]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Expand the next cell to find the solution to 2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    },
    "tags": []
   },
   "source": [
    "If $X$ and $A$ are sampled from the uniform distribution over the state-action space. Then,\n",
    "\n",
    "$$\n",
    "\\nabla_\\mu \\mathcal{L}(v, \\mu) = |\\S||\\A|\\cdot \\mathbb{E}[ r(X,A) + \\gamma v(X') - v(X) ]\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EjncDj_f16rZ",
    "tags": []
   },
   "outputs": [],
   "source": [
    "SampleDT = [(x, int) for x in [\"s0\", \"s\", \"a\", \"r\", \"next_s\"]]\n",
    "\n",
    "\n",
    "def collect_sample(env, s, a):\n",
    "    s0 = env.s\n",
    "    \n",
    "    # A comment in memory of the hours spent before\n",
    "    # realizing I had to put `unwrapped` here.\n",
    "    env.unwrapped.s = s \n",
    "    next_s, r, *_ = env.step(a)\n",
    "    \n",
    "    return np.rec.array((s0, s, a, r, next_s), dtype=SampleDT)\n",
    "\n",
    "\n",
    "class GenerativePD(DistribPD):\n",
    "\n",
    "    def grad_mu(self, v):\n",
    "        # compute an unbiased estimate of the gradient of the\n",
    "        # Lagrangian wrt to mu evaluated at point `v`,\n",
    "        # using the sample `self.sample`\n",
    "        w = self.sample\n",
    "\n",
    "        grad = np.zeros((self.S, self.A))\n",
    "        \n",
    "        # use the sample w to assign the correct components of the gradient\n",
    "        \n",
    "        return grad\n",
    "\n",
    "    def grad_v(self, mu):\n",
    "        # compute an unbiased estimate of the gradient of the\n",
    "        # Lagrangian wrt to v evaluated at point `mu`,\n",
    "        # using the sample `self.sample`\n",
    "        w = self.sample\n",
    "\n",
    "        grad = np.zeros(self.S)\n",
    "        \n",
    "        # use the sample w to assign the correct components of the gradient\n",
    "\n",
    "        return grad\n",
    "\n",
    "\n",
    "    def step(self, t):  \n",
    "        # update for mu\n",
    "        \n",
    "        # sample a state and action from the uniform distribution\n",
    "        # Hint: use the generator self.rng\n",
    "        s = ...\n",
    "        a = ...\n",
    "        \n",
    "        # collect a new sample of the reward and next state given s, a\n",
    "        self.sample = collect_sample(self.env, s, a)\n",
    "        \n",
    "        # compute a step of mu\n",
    "        mu = self.step_mu()\n",
    "        \n",
    "        # reset the env to make sure that the next time we call\n",
    "        # `collect_sample` we will also get a fresh sample from\n",
    "        # the initial state distribution nu0\n",
    "        self.env.reset()\n",
    "        \n",
    "        # update for v\n",
    "        \n",
    "        # sample a state and action from `self.mu`\n",
    "        s = self.sample_state()\n",
    "        a = self.sample_action(s)\n",
    "        \n",
    "        # get a new sample of (s0, r, s')\n",
    "        self.sample = collect_sample(self.env, s, a)\n",
    "\n",
    "        # compute a step of v\n",
    "        v = self.step_v()\n",
    "        \n",
    "        # Finally, set the variables\n",
    "        self.v = v\n",
    "        self.mu = mu\n",
    "        \n",
    "        # make sure we store the cumulative mu, to compute our output policy\n",
    "        self.mu_sum += mu\n",
    "\n",
    "    def sample_state(self):\n",
    "        \"Sample a state from self.mu\"\n",
    "        # hint: use self.rng (and self.mu)\n",
    "        return ...\n",
    "    \n",
    "    def sample_action(self, state):\n",
    "        \"Sample an action from self.mu, for the supplied `state`\"\n",
    "        p = self.mu[state]\n",
    "        return ...\n",
    "\n",
    "\n",
    "class LoggingGenerativePD(LoggingMixin, GenerativePD):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's run our algorithm!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0J8rdlRpbk2d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "n = 40_000\n",
    "lr = 0.0031\n",
    "\n",
    "agent = LoggingGenerativePD(\n",
    "    env,\n",
    "    lr_v = lr,\n",
    "    lr_mu = lr / 3,\n",
    "    log_every = 1000,\n",
    "    seed = ALG_SEED\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "f_DlIrbPUEaN",
    "tags": []
   },
   "outputs": [],
   "source": [
    "for t in range(1, n):\n",
    "    #agent.lr = 1/np.sqrt(t)\n",
    "    agent.step(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 452
    },
    "id": "9rrEzAWRULaN",
    "outputId": "be900d93-e4e5-47ba-faad-53142ede1bef",
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.plot(agent._rewards)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 440
    },
    "id": "SCoWFdYfm_A8",
    "outputId": "9914aad1-900b-41f1-bf4e-2252e3187cc6",
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_policy(agent.policy)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}