dp03_bandits-and-gittins-index.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dd0HlxDDOJ93"
   },
   "source": [
    "# <center>Multi-armed bandits and Gittins index</center>\n",
    "### <center>Alfred Galichon (NYU & Sciences Po)</center>\n",
    "## <center>'math+econ+code' masterclass series</center>\n",
    "#### <center>With python code examples</center>\n",
    "© 2018–2023 by Alfred Galichon. Past and present support from NSF grant DMS-1716489, ERC grant CoG-866274 are acknowledged, as well as inputs from contributors listed [here](http://www.math-econ-code.org/team).\n",
    "\n",
    "**If you reuse material from this masterclass, please cite as:**<br>\n",
    "Alfred Galichon, 'math+econ+code' masterclass series. https://www.math-econ-code.org/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kTSWIALlOJ98"
   },
   "source": [
    "# Introduction\n",
    "\n",
    "### Learning objectives\n",
    "\n",
    "* Gittins index\n",
    "* LCP formulation\n",
    "\n",
    "\n",
    "### References\n",
    "\n",
    "* Bertsekas TBC.\n",
    "* Weber (2016). *Optimization and Control*. Lecture notes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bKgWhZyKOJ98"
   },
   "source": [
    "\n",
    "# Setup\n",
    "\n",
    "Consider a finite state space $\\mathcal{X}$, and an action space $\\mathcal{Y}\n",
    "$. The short-term reward associated with taking action $y$ in space $x$ is $%\n",
    "\\Phi _{xy}$. The decision variable $\\mu _{xy}\\geq 0$ is the number of units of type $x$ for\n",
    "which action $y$ is taken. \n",
    "\n",
    "If action $y$ is taken for a unit in state $x$, then the probability of\n",
    "transitioning to a state $x^{\\prime }$ is $N_{x^{\\prime },x}^{y}$. We have \n",
    "$$\n",
    "N=\\sum_{y}N^{y}\\left( e^{y}\\right) ^{\\top },\n",
    "$$\n",
    "so that $N_{x^{\\prime },xy}=N_{x^{\\prime },x}^{y}$.\n",
    "\n",
    "Consider $M=\\mathbf{I}_{\\mathcal{X}}\\otimes \\mathbf{1}_{\\mathcal{Y}}^{\\top }$, so that $M_{x^{\\prime },xy}=1\\left\\{ x=x^{\\prime }\\right\\} $. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "HkHtqzqCOJ9-"
   },
   "outputs": [],
   "source": [
    "#!pip install gurobipy\n",
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "import gurobipy as grb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bandits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each bandit $i \\in \\mathcal I$ has a state space $\\mathcal Z^i$. The state space is endowed with a markov chain $P^i_{z^\\prime, z}$ which is the probability of a transition from $z$ to $z^\\prime$. A reward $\\phi^i_z$ is associated with bandit $i$ accessing state $z$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "rQIcaX8uU2p0"
   },
   "outputs": [],
   "source": [
    "class Bandit():\n",
    "    def __init__(self,P_zp_z,phi_z, beta):\n",
    "        self.nbz = len(phi_z)\n",
    "        self.P_zp_z = P_zp_z\n",
    "        self.phi_z = phi_z\n",
    "        self.beta = beta\n",
    "        \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two-armed bandits\n",
    "\n",
    "In a single-armed bandit, one has the choice between continuing and stopping. Thus the action space $\\mathcal Y = \\{0,1 \\}$ consists of continuing or stopping and getting $g$ forever.\n",
    "\n",
    "For each $z^{\\ast }\\in \\mathcal{Z}$, the Gittins index $g_{z^{\\ast }}$ at state $z^{\\ast }$ is the solution to the following equation\n",
    "$$\n",
    "\\left\\{ \n",
    "\\begin{array}{l}\n",
    "p_{z}=\\max \\left\\{ \\frac{g_{z^{\\ast }}}{1-\\beta },\\phi _{z}+\\beta \\left(\n",
    "P^{\\top }p\\right) _{z}\\right\\} \\forall z\\in \\mathcal{Z} \\\\ \n",
    "\\frac{g_{z^{\\ast }}}{1-\\beta }=\\phi _{z^{\\ast }}+\\beta \\left( P^{\\top\n",
    "}p\\right) _{z^{\\ast }}%\n",
    "\\end{array}%\n",
    "\\right. \n",
    "$$\n",
    "which can be solved using the following LCP \n",
    "$$\n",
    "\\left\\{\n",
    "\\begin{array}{l}\n",
    "0\\leq \\rho _{z}^{1}\\perp \\mu _{z1}\\geq 0 \\\\ \n",
    "0\\leq \\rho _{z}^{2}\\perp \\mu _{z2}\\geq 0 \\\\ \n",
    "\\mu _{z1}+\\mu _{z2}=1 \\\\ \n",
    "\\rho _{z}^{1}=p_{z}-\\frac{g }{1-\\beta } \\\\ \n",
    "\\rho _{z}^{2}=p_{z}-\\phi _{z}-\\beta \\left( P^{\\top }p\\right) _{z} \\\\ \n",
    "\\frac{g }{1-\\beta }=\\phi _{z^{\\ast }}+\\beta \\left( P^{\\top }p\\right)\n",
    "_{z^{\\ast }}.\n",
    "\\end{array}\n",
    "\\right. \n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.3"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "abs(-1.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Bandit_gittins_lcp(self):\n",
    "    self.g_z = np.zeros(self.nbz)\n",
    "    for zstar in range(self.nbz):\n",
    "        m=grb.Model()\n",
    "        m.Params.OutputFlag = 0\n",
    "        m.Params.NonConvex = 2\n",
    "        mu1_z = m.addMVar(self.nbz) \n",
    "        mu2_z = m.addMVar(self.nbz) \n",
    "        rho1_z = m.addMVar(self.nbz) \n",
    "        rho2_z = m.addMVar(self.nbz) \n",
    "        p_z = m.addMVar(self.nbz,lb = - grb.GRB.INFINITY) \n",
    "        g = m.addMVar(1,lb = - grb.GRB.INFINITY)\n",
    "        m.setObjective(mu1_z@ rho1_z + mu2_z@ rho2_z, grb.GRB.MINIMIZE)\n",
    "        m.addConstr(mu1_z+mu2_z == np.ones(self.nbz))\n",
    "        m.addConstr(rho1_z == p_z - g * np.ones(self.nbz) / (1-self.beta))\n",
    "        m.addConstr(rho2_z == p_z - self.phi_z - self.beta * self.P_zp_z.T @ p_z )\n",
    "        m.addConstr(g / (1-self.beta) == self.phi_z[zstar] +  self.beta * (self.P_zp_z.T @ p_z)[zstar] )\n",
    "        m.optimize()\n",
    "        self.g_z[zstar] = g.X\n",
    "    return(self.g_z)\n",
    "    \n",
    "Bandit.gittins_lcp = Bandit_gittins_lcp\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: single machine scheduling\n",
    "\n",
    "A single machine, $n$ jobs. Job $i$ takes time $t_{i}$ to process and reward $r_{i}$ is obtained at the end of\n",
    "the job. \n",
    "\n",
    "Stopping time for job $i$=$t_{i}$. Hence Gittins index associated with job $i$ is\n",
    "\n",
    "$\\frac{r_{i}\\beta ^{t_{i}}}{1+\\beta +...+\\beta ^{t_{i}}}=\\frac{r_{i}\\beta\n",
    "^{t_{i}}\\left( 1-\\beta \\right) }{1-\\beta ^{t_{i}}}.$\n",
    "\n",
    "Therefore jobs should be processed in the decreasing order of the Gittins\n",
    "index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_machine_scheduling_bandit(nbperiods,payoff,beta,verbose = True):\n",
    "    nbz=nbperiods\n",
    "    P_zp_z = np.diag(np.ones(nbz-1),-1)\n",
    "    P_zp_z[-1,-1]=1\n",
    "    phi_z = np.zeros(nbz)\n",
    "    phi_z[-2] = payoff\n",
    "    if verbose:\n",
    "        print('P_zp_z=\\n',P_zp_z)\n",
    "        print('phi_z=\\n',phi_z)\n",
    "    return(Bandit(P_zp_z,phi_z,beta))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution via LCP"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We solve for the Gittins indices using our LCP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P_zp_z=\n",
      " [[0. 0. 0. 0. 0. 0.]\n",
      " [1. 0. 0. 0. 0. 0.]\n",
      " [0. 1. 0. 0. 0. 0.]\n",
      " [0. 0. 1. 0. 0. 0.]\n",
      " [0. 0. 0. 1. 0. 0.]\n",
      " [0. 0. 0. 0. 1. 1.]]\n",
      "phi_z=\n",
      " [ 0.  0.  0.  0. 35.  0.]\n",
      "Set parameter Username\n",
      "Academic license - for non-commercial use only - expires 2023-12-23\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([ 6.30090993,  8.08871593, 11.07361963, 17.05128205, 35.        ])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "the_payoff = 35\n",
    "the_time = 6\n",
    "the_beta = 0.95\n",
    "the_bandit = single_machine_scheduling_bandit(the_time,the_payoff,the_beta)\n",
    "the_bandit.gittins_lcp()[:-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution via Jacobi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Bandit_gittins_jacobi(self, tol = 1e-5, maxit = 1000):\n",
    "    self.g_z = np.zeros(self.nbz)\n",
    "    for zstar in range(self.nbz):\n",
    "        p_z = (self.phi_z.min()  / (1-self.beta) ) * np.ones(self.nbz) \n",
    "        g = (1 - self.beta)* self.phi_z.min() + self.beta * (self.phi_z.min()  / (1-self.beta) )\n",
    "        cont,nbit = True,0\n",
    "        while cont:\n",
    "            PTp_z = self.P_zp_z.T @ p_z\n",
    "            newp_z = np.maximum( g / (1-self.beta) , self.phi_z+self.beta * PTp_z)\n",
    "            newg = (1-self.beta)*(self.phi_z[zstar]+self.beta*PTp_z[zstar])\n",
    "            if (max(np.linalg.norm(p_z - newp_z, ord=1),abs(newg-g))<tol) or (nbit > maxit):\n",
    "                cont = False\n",
    "            nbit += 1\n",
    "            p_z,g = newp_z,newg\n",
    "        self.g_z[zstar] = g\n",
    "    return(self.g_z)\n",
    "\n",
    "Bandit.gittins_jacobi=Bandit_gittins_jacobi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 6.30090824,  8.0887134 , 11.07361651, 17.05127724, 34.99999061])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Bandit.gittins_jacobi(the_bandit)[:-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution in closed form"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Actually, the problem can be solved in closed form. We can compare the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 6.30090993,  8.08871593, 11.07361963, 17.05128205, 35.        ])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "the_times = np.arange(the_time-2,-1,-1)\n",
    "the_payoff*(the_beta**the_times)*(1-the_beta)/(1-the_beta**(the_times+1)) \n",
    "# careful: formula in Weber's lecture notes has a typo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-armed bandits\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Bandits():\n",
    "    def __init__(self,bandits):\n",
    "        self.bandits = bandits\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "the_bandits = Bandits([single_machine_scheduling_bandit(3,4,0.9,False),single_machine_scheduling_bandit(2,6,0.9,False)])\n"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}