dp02_tbc_markov-decision-processes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dd0HlxDDOJ93"
   },
   "source": [
    "# <center>Markov decision processes</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",
    "* Dynamic programming\n",
    "* Stationarity\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",
    "\n",
    "### Dynamic programming\n",
    "\n",
    "If $(q_x)_x$ is the present period distribution over states, and if $(p^\\prime_x)_x$ is the next period payoff associated to each state, we can compute the intertemporal value as:\n",
    "$W\\left( q,p^{\\prime }\\right) =\\max_{\\mu \\geq 0}\\left\\{ \\mu ^{\\top }\\left(\n",
    "\\Phi +\\beta N^{\\top }p^{\\prime }\\right) :M\\mu =q\\right\\} $.\n",
    "\n",
    "This induces a policy $\\pi_{y|x} = \\mu_{xy}/q_x$.\n",
    "\n",
    "Let $p_{x}$ be the present-period reward associated with state $x$, and let  $q_{x}^{\\prime }$ be the next-period number of units in state $x$. We have<div> \n",
    "$p \\in \\partial_q W(q,p^\\prime)$<div> \n",
    "where $\\partial_q$ denotes the superdifferential in $q$, and<div>\n",
    "$q^\\prime \\in \\partial_p W(q,p^\\prime)$<div> \n",
    "where $\\partial_p$ denotes the subdiffernential in $p$."
   ]
  },
  {
   "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": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mdp():\n",
    "    def __init__(self,list_Phi_x,list_N_xp_x,beta, pp_x=None,q_x=None):\n",
    "        self.nbx = len(list_Phi_x[0])\n",
    "        self.nby = len(list_Phi_x)\n",
    "        self.Phi_xy = np.block( [Phi_x.reshape((-1,1)) for Phi_x in list_Phi_x]).flatten()\n",
    "        self.N_xp_xy = np.block( [N_xp_x.reshape((-1,1)) for N_xp_x in list_N_xp_x]  ).reshape(self.nbx,-1)\n",
    "        self.beta = beta\n",
    "        self.pp_x = pp_x\n",
    "        self.q_x = q_x\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eazveEm1OJ-G"
   },
   "source": [
    "Define the model and call the solver by:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling a MDP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Mdp_sampler_init(self,state=0):\n",
    "    self.state = state\n",
    "    self.payoff = 0.0\n",
    "    self.period = 0\n",
    "\n",
    "Mdp.sampler_init = Mdp_sampler_init\n",
    "\n",
    "def Mdp_sampler_forward(self,y,seed=None,verbose=False):\n",
    "    if seed is not None:\n",
    "        np.random.seed(seed)\n",
    "    xy = self.nby * self.state + y\n",
    "    oldstate = self.state\n",
    "    self.state = np.random.choice(list(range(self.nbx)),1, p=self.N_xp_xy[:,xy])[0]\n",
    "    self.period += 1\n",
    "    self.payoff = self.payoff/self.beta + self.Phi_xy[xy]  \n",
    "    if verbose:\n",
    "        print('transition',str(oldstate),'->',str(self.state))\n",
    "        print('new state',self.state)\n",
    "        print('new period',self.period)\n",
    "        print('new payoff',self.payoff)\n",
    "    return(self.state)\n",
    "    \n",
    "Mdp.sampler_forward = Mdp_sampler_forward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_simple_mdp(nbperiods,payoff1,payoff2,beta,verbose = True):\n",
    "    nbz=nbperiods\n",
    "    P1_zp_z = np.diag(np.ones(nbz-1),-1)\n",
    "    P1_zp_z[-1,-1]=1\n",
    "    P2_zp_z = np.zeros((nbz,nbz))\n",
    "    P2_zp_z[0,:]=1\n",
    "    phi1_z = np.zeros(nbz)\n",
    "    phi1_z[-2] = payoff1\n",
    "    phi2_z = np.zeros(nbz)\n",
    "    phi2_z[0] = payoff2\n",
    "    if verbose:\n",
    "        print('P1_zp_z=\\n',P1_zp_z)\n",
    "        print('P2_zp_z=\\n',P2_zp_z)\n",
    "        print('phi1_z=\\n',phi1_z)\n",
    "        print('phi2_z=\\n',phi2_z)\n",
    "    return(Mdp([phi1_z,phi2_z],[P1_zp_z,P2_zp_z],beta))\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "P1_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",
      "P2_zp_z=\n",
      " [[1. 1. 1. 1. 1. 1.]\n",
      " [0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0.]]\n",
      "phi1_z=\n",
      " [0. 0. 0. 0. 5. 0.]\n",
      "phi2_z=\n",
      " [-4.  0.  0.  0.  0.  0.]\n"
     ]
    }
   ],
   "source": [
    "the_mdp = build_simple_mdp(6,5,-4,0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "the_mdp.sampler_init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "transition 0 -> 1\n",
      "new state 1\n",
      "new period 1\n",
      "new payoff 0.0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "the_mdp.sampler_forward(0,verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "transition 1 -> 0\n",
      "new state 0\n",
      "new period 2\n",
      "new payoff 0.0\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "the_mdp.sampler_forward(1,verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6, 12)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "the_mdp.N_xp_xy.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Mdp_simulate(self,nbperiod,pi_x_y,seed = None,initial_state=0,verbose = 0):\n",
    "    if seed is not None:\n",
    "        np.random.seed(seed)\n",
    "    self.sampler_init(initial_state)\n",
    "    for t in range(nbperiod):\n",
    "        y =  np.random.choice(list(range(self.nby)),1, p=pi_x_y[self.state,:])[0]\n",
    "        if verbose>0:\n",
    "            print(y)\n",
    "        self.sampler_forward(y)\n",
    "    return(self.payoff * (self.beta**self.period) )\n",
    "\n",
    "Mdp.simulate=Mdp_simulate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Value of a policy\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assume I know initial value $q$ and policy $\\mu $. What is the total value?\n",
    "\n",
    "\n",
    "If $\\pi _{y|x}$ is the conditional distribution, we have\n",
    "\n",
    "$\\mu _{xy}=\\pi _{y|x}q_{x}$\n",
    "\n",
    "$q_{x^{\\prime }}^{\\prime }=\\sum_{x}\\sum_{y}N_{x^{\\prime }|xy}\\pi _{y|x}q_{x}$\n",
    "\n",
    "we can define $\\tilde{N}_{x^{\\prime }|x}=\\sum_{y}N_{x^{\\prime }|xy}\\pi\n",
    "_{y|x} $\n",
    "\n",
    "and we have $q_{x^{\\prime }}^{t}=\\left( \\tilde{N}\\right) ^{t}q^{0}$\n",
    "\n",
    "and the surplus is equal to\n",
    "\n",
    "$\\sum_{t}\\sum_{x}q_{x}^{t}\\beta ^{t}{\\tilde \\phi} _{x}=\\sum_{t=0}^{\\infty }\\tilde{\\phi%\n",
    "}^{\\top }\\left( \\beta ^{t+1}\\tilde{N}^{t}\\right) q^{0}$\n",
    "\n",
    "where $\\tilde{\\phi}_{x}=\\sum_{y}\\pi _{y|x}\\Phi _{xy}$\n",
    "\n",
    "thus ${\\tilde \\phi} ^{\\top }\\beta \\sum_{t=0}^{\\infty }\\left( \\beta \\tilde{N}\\right)\n",
    "^{t}q^{0}$\n",
    "\n",
    "that is $\\tilde{\\phi}^{\\top }\\beta \\left( 1-\\beta \\tilde{N}\\right) ^{-1}q^{0}\n",
    "$\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Mdp_policy_value(self,pi_x_y):\n",
    "    Ntilde_xp_x = (pi_x_y[None,:,:] * self.N_xp_xy.reshape((self.nbx,self.nbx,-1))).sum(axis=2)\n",
    "    phithilde_x = (pi_x_y * self.Phi_xy.reshape( (self.nbx,-1))).sum(axis=1)\n",
    "    p_x = the_mdp.beta*np.linalg.solve((np.eye(the_mdp.nbx) - the_mdp.beta * Ntilde_xp_x.T),phithilde_x)\n",
    "    return(p_x)\n",
    "    \n",
    "Mdp.policy_value = Mdp_policy_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([6.3011275 , 7.00125278, 7.77916975, 8.64352195, 9.60391327,\n",
       "       5.67101475])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pi_x_y=np.zeros((the_mdp.nbx,the_mdp.nby))\n",
    "pi_x_y[0:-1,0]=1\n",
    "pi_x_y[-1,1]=1\n",
    "\n",
    "the_mdp.policy_value(pi_x_y)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([6.3011275 , 7.00125278, 7.77916975, 8.64352195, 9.60391327,\n",
       "       5.67101475])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array([the_mdp.simulate(1000,pi_x_y,initial_state=x) for x in range(the_mdp.nbx)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving for the optimal policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Mdp_solve_lp(self):\n",
    "    m=grb.Model()\n",
    "    m.Params.NonConvex = 2\n",
    "    mu_xy = m.addMVar(self.nbx*self.nby)\n",
    "    rho_xy = m.addMVar(self.nbx*self.nby)\n",
    "    p_x = m.addMVar(self.nbx, lb = - grb.GRB.INFINITY)\n",
    "    qp_x = m.addMVar(self.nbx )\n",
    "    m.setObjective(mu_xy @ rho_xy , grb.GRB.MINIMIZE)\n",
    "    if self.pp_x is not None:\n",
    "        m.addConstr(rho_xy+self.Phi_xy+self.beta*self.N_xp_xy.T @ self.pp_x - self.M_xp_xy.T @ p_x == 0 )\n",
    "    else:\n",
    "        m.addConstr(rho_xy+self.Phi_xy + (self.beta*self.N_xp_xy - self.M_xp_xy).T @ p_x == 0 )\n",
    "    m.addConstr(self.N_xp_xy @ mu_xy == qp_x)\n",
    "    if self.q_x is not None:\n",
    "        m.addConstr(self.M_xp_xy @ mu_xy == self.q_x)\n",
    "    else:\n",
    "        m.addConstr(self.M_xp_xy @ mu_xy == qp_x)\n",
    "        m.addConstr(qp_x.sum()==1)\n",
    "        \n",
    "    m.optimize()\n",
    "    return(p_x.X,qp_x.X,mu_xy.X)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "We have the following orthogonality conditions relating $p,p^\\prime,q$ and $q^\\prime$\n",
    "$$\\left\\{ \n",
    "\\begin{array}{l}\n",
    "0\\leq \\Phi +\\beta N^{\\top }p^{\\prime }-M^{\\top }p\\perp \\mu \\geq 0 \\\\ \n",
    "M\\mu =q \\\\ \n",
    "N\\mu =q^{\\prime }%\n",
    "\\end{array}%\n",
    "\\right. \n",
    "$$\n",
    "which we can rewrite as\n",
    "$$\n",
    "\\left\\{ \n",
    "\\begin{array}{l}\n",
    "0\\leq \\rho \\perp \\mu \\geq 0 \\\\ \n",
    "\\rho =M^{\\top }p-\\beta N^{\\top }p^{\\prime }-\\Phi  \\\\ \n",
    "M\\mu =q \\\\ \n",
    "N\\mu =q^{\\prime }%\n",
    "\\end{array}%\n",
    "\\right. \n",
    "$$\n",
    "\n",
    "Note that $p=p^{\\prime }$ and $q=q^{\\prime }$ yields the stationary solution.\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
}