Policy Learning with Observational Data Tutorial (2).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "37f61be5",
   "metadata": {},
   "source": [
    "<center>\n",
    "    <h1>Policy Learning with Observational Data: <i>A Tutorial</i></h1>\n",
    "    <h2>Implementing Susan Athey and Stefan Wager (2021)</h2>\n",
    "    <h3>Allison Towey, <i>The University of Chicago</i></h3>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c06e9a71",
   "metadata": {},
   "source": [
    "### Table of Contents\n",
    "\n",
    "- [1. Introduction](#section-1-introduction)\n",
    "- [2. Getting Started](#section-2)\n",
    "- [3. ***Step 1***: Estimation of Nuisance Components ](#section-3)\n",
    "- [4. ***Step 2***: Computing Doubly Robust Scores](#section-4)\n",
    "- [5. ***Step 3*** Fitting the Policy Tree](#section-5)\n",
    "- [6. Evaluation](#section-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76647874",
   "metadata": {},
   "source": [
    "# 1. Introduction <a id=\"section-1-introduction\"></a>\n",
    "\n",
    "This tutorial explores the implementation of Athey and Wager's framework from **\"Policy Learning with Observational Data\"** (2021), which addresses the challenge of learning treatment assignment policies while considering feasibility constraints. It provides step-by-step guidance using a real dataset to facilitate the implementation of the framework in Python. The goal is to provide readers with a solid understanding of effective tools to apply policy learning methodology.\n",
    "\n",
    "<h2>What is policy learning?</h2>\n",
    "\n",
    "Policy learning plays a crucial role in enabling adaptive, optimal decision-making in complex and uncertain environments, with applications spanning various domains and disciplines. Accurately and efficiently learning optimal treatment assignments, which involves mapping individual characteristics to treatment assignments, is important for policy decision-making and a common challenge in applied economics and statistics.\n",
    "\n",
    "Across various domains, practitioners seek to utilize observational data to learn treatment assignment policies that adhere to specific constraints, as the treatment assignment problem is seldom encountered in an unconstrained environment. These policies may be subject to constraints such as budget limitations, fairness considerations, simplicity requirements, or other functional requirements, which can limit the range of feasible treatment assignments. In real life applications, there may be several reasons to require that the policy belong to a restricted class of policies, such as:\n",
    "<ul>\n",
    "<li>Some variables such as gender or race may not be able to be used to determine the treatment allocation (ethically, or legally).</li><li> Interpretability and functional form simplicity could be important if understandability and explainability is fundamental to policy adoption and auditing.</li><li>Technical efficiency may be necessary or beneficial in order to achieve optimal performance or resource utilization.</li></ul>\n",
    "    \n",
    "Practitioners aim to find an optimal set of treatment assignments while considering these feasibility constraints. The objective is to minimize regret, which represents the difference between the utility of the chosen policy and the highest utility achievable among all policies within the feasible set."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d77dd3d0",
   "metadata": {},
   "source": [
    "<h2> Our objective </h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39a91776",
   "metadata": {},
   "source": [
    "1. We would like to select a function $\\pi$ that maps observed characteristics to an available treatment. This function is called a <strong>policy</strong>.\n",
    "\n",
    "\n",
    "2. Using observational data, we want to learn a policy $\\pi$ that maps a subject’s **characteristics**, $X_i \\in X$, to a **treatment decision** $\\pi: X \\to ${$0, 1$}.\n",
    "\n",
    "\n",
    "3. This learned policy (𝛑̂) must satisfy certain **given constraints** and belong to the set of feasible policies ($\\Pi$), i.e.: 𝛑̂ ∈ $\\Pi$\n",
    "\n",
    "\n",
    "4. We seek guarantees on the **regret**: R(𝛑̂), or the difference between the expected utility of the learned policy 𝛑̂ and the best utility of any policy in the set of feasible policies ($\\Pi$)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a81148",
   "metadata": {},
   "source": [
    "<h2>So, how do we do this?</h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d267d1a",
   "metadata": {},
   "source": [
    "Athey and Wager's <strong>\"Policy Learning with Observational Data\" </strong> (published 2021 in Econometrica) develops an algorithm that uses observational data to determine the best way to assign treatments to individuals given a set of constraints. Informed by the theory of semi-parametrically efficient estimation, they provide assurances that their method will result in minimal regret. \n",
    "\n",
    "Their method seeks to learn treatment assignment policies that can: \n",
    "<ol>\n",
    "    <li> handle specified constraints and </li>\n",
    "    <li> use observational (i.e. not only Random Controlled Trial) data. </li>\n",
    "   </ol>\n",
    "The purpose of this tutorial is to demonstrate the application of their method to a real dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45affff9",
   "metadata": {},
   "source": [
    "<h2>Setting Up Our Framework and Assumptions</h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "262bdb19",
   "metadata": {},
   "source": [
    "### Causal Framework\n",
    "\n",
    "Before we begin the tutorial, we set up our causal framework.\n",
    "\n",
    "Each unit in our data set will be represented by:\n",
    "\n",
    "- $X_i$: observable pre-treatment characteristics\n",
    "- $W_i ∈{0,1}$ whether the unit was treated (1) or not treated (0)\n",
    "- $Y_i^{obs} ∈ ℝ$: the observed binary outcome for that unit\n",
    "\n",
    "The **law of total expectation**, also known as the law of iterated expectation, states that the expected value of a random variable can be calculated by taking the conditional expectation with respect to another random variable and then averaging over all possible values of that variable. \n",
    "\n",
    "This law is expressed as:\n",
    "<center>$E[X] = E[E[X|Y]]$</center>\n",
    "\n",
    "We additionally adopt the **potential outcome framework** introduced by Rubin (1974) to study causal relationships. This framework considers the existence of two potential outcomes for each unit: one under the treatment condition ($Y_i(1)$) and one under the control condition ($Y_i(0)$). However, in practice, we can only observe a single outcome for each individual based on whether they receive the treatment or control. This creates a \"missing value\" problem, where we are interested in analyzing counterfactual scenarios. We have underlying random variables including the treatment assignment ($X_i$), the potential outcome under control ($Y_i(0)$), and the potential outcome under treatment ($Y_i(1)$). But we can only observe either ($X_i, Y_i(0)$) for control units or ($X_i, Y_i(1)$) for treated units. Therefore, causal inference involves addressing this missing outcome issue to make inferences and estimate treatment effects or intervention outcomes. \n",
    "\n",
    "Consequently, the observed outcome can be expressed as:\n",
    "<center>$Y_i^{obs}=W_iY_i(1)+(1−W_i)Y_i(0)$ </center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a52ecc01",
   "metadata": {},
   "source": [
    "### Assumptions\n",
    "Additionally, we will be making two identification assumptions that will allow us to use the methods below.\n",
    "\n",
    "#### Assumption 1: Unconfoundedness/Conditional Independence \n",
    "The unconfoundedness assumption states that the treatment assignment is independent of potential outcomes given the observed covariates or confounders. Once we account for the observable characteristics or variables that might affect both the treatment assignment and the outcome, the unconfoundedness assumption implies that there are no hidden or unobserved factors that influence both the treatment assignment and the outcome. Formally, \n",
    "<center>$Y_i(1),Y_i(0)⊥ W_i | X_i$</center>\n",
    "\n",
    "#### Assumption 2: Positivity/Overlap\n",
    "The positivity assumption states that for any combination of observed covariates or confounders, there is a non-zero probability of receiving both the treatment and the control condition.It therefore implies that there are no regions in the covariate space where all individuals are either always treated or always untreated. It ensures that there is sufficient overlap in the distribution of covariates between the treated and untreated groups, allowing for meaningful comparisons and estimation of treatment effects. Formally, \n",
    "<center>$ 0<P (W=1 | X=x)<1$</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d3d4e35",
   "metadata": {},
   "source": [
    "# 2. Getting Started<a id=\"section-2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77355e84",
   "metadata": {},
   "source": [
    "To demonstrate how to implement algorithm, let's first set up our coding environment and data.\n",
    "\n",
    "This tutorial will use the coding language Python, along with several important libraries imported below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4114293b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Turn off all warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Import libraries\n",
    "# If any of these libraries are not installed, you may do so by running the command: !pip install your_library_name\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import random\n",
    "from sklearn.model_selection import train_test_split\n",
    "from IPython.display import SVG, display\n",
    "import rpy2.robjects as robjects\n",
    "from rpy2.robjects import pandas2ri\n",
    "\n",
    "# We also want to set a seed for replicability\n",
    "random.seed(123456)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b2e935f",
   "metadata": {},
   "source": [
    "We will be working with the 'welfare' dataset from the GSB Social Impact Lab, originally from “Modeling Heterogeneous Treatment Effects in Survey Experiments with Bayesian Additive Regression Trees” (Green and Kern 2012)\n",
    "\n",
    "It is additionally possible to run this code with other pre-cleaned datasets from the GSB Social Impact Lab such as the social dataset from “Social Pressure and Voter Turnout” (Gerber, Green, and Larimer, 2008) and the charitable dataset from “Does Price Matter in Charitable Giving?” (Karlan and List, 2007).\n",
    "\n",
    "The data provided by the GSB Social Impact Lab is largely pre-cleaned, so there are relatively few data-cleaning or generating processes needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "06d1c0c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in our data\n",
    "df = pd.read_csv('https://raw.githubusercontent.com/gsbDBI/ExperimentData/master/Welfare/ProcessedData/welfarenolabel3.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec9c53c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['year',\n",
       " 'id',\n",
       " 'wrkstat',\n",
       " 'hrs1',\n",
       " 'hrs2',\n",
       " 'evwork',\n",
       " 'occ',\n",
       " 'prestige',\n",
       " 'wrkslf',\n",
       " 'wrkgovt',\n",
       " 'commute',\n",
       " 'occ80',\n",
       " 'prestg80',\n",
       " 'indus80',\n",
       " 'marital',\n",
       " 'agewed',\n",
       " 'divorce',\n",
       " 'widowed',\n",
       " 'spwrksta',\n",
       " 'sphrs1',\n",
       " 'sphrs2',\n",
       " 'spevwork',\n",
       " 'spocc80',\n",
       " 'sppres80',\n",
       " 'spind80',\n",
       " 'sibs',\n",
       " 'childs',\n",
       " 'age',\n",
       " 'agekdbrn',\n",
       " 'educ',\n",
       " 'paeduc',\n",
       " 'maeduc',\n",
       " 'speduc',\n",
       " 'degree',\n",
       " 'padeg',\n",
       " 'madeg',\n",
       " 'spdeg',\n",
       " 'sex',\n",
       " 'race',\n",
       " 'res16',\n",
       " 'reg16',\n",
       " 'mobile16',\n",
       " 'family16',\n",
       " 'mawork',\n",
       " 'mawkborn',\n",
       " 'born',\n",
       " 'parborn',\n",
       " 'granborn',\n",
       " 'hompop',\n",
       " 'babies',\n",
       " 'preteen',\n",
       " 'teens',\n",
       " 'adults',\n",
       " 'unrelat',\n",
       " 'earnrs',\n",
       " 'income',\n",
       " 'rincome',\n",
       " 'income86',\n",
       " 'partyid',\n",
       " 'polviews',\n",
       " 'w',\n",
       " 'y',\n",
       " '_merge',\n",
       " 'wrkstat_miss',\n",
       " 'hrs1_miss',\n",
       " 'hrs2_miss',\n",
       " 'evwork_miss',\n",
       " 'wrkslf_miss',\n",
       " 'wrkgovt_miss',\n",
       " 'commute_miss',\n",
       " 'occ80_miss',\n",
       " 'indus80_miss',\n",
       " 'agewed_miss',\n",
       " 'divorce_miss',\n",
       " 'widowed_miss',\n",
       " 'spwrksta_miss',\n",
       " 'sphrs1_miss',\n",
       " 'sphrs2_miss',\n",
       " 'spevwork_miss',\n",
       " 'spocc80_miss',\n",
       " 'spind80_miss',\n",
       " 'sibs_miss',\n",
       " 'agekdbrn_miss',\n",
       " 'paeduc_miss',\n",
       " 'maeduc_miss',\n",
       " 'speduc_miss',\n",
       " 'padeg_miss',\n",
       " 'madeg_miss',\n",
       " 'spdeg_miss',\n",
       " 'res16_miss',\n",
       " 'mobile16_miss',\n",
       " 'family16_miss',\n",
       " 'mawork_miss',\n",
       " 'mawkborn_miss',\n",
       " 'born_miss',\n",
       " 'parborn_miss',\n",
       " 'granborn_miss',\n",
       " 'unrelat_miss',\n",
       " 'rincome_miss',\n",
       " 'income86_miss',\n",
       " 'polviews_miss',\n",
       " 'polviews_num',\n",
       " 'income_num',\n",
       " 'racdif1',\n",
       " 'racdif2',\n",
       " 'racdif3',\n",
       " 'racdif4',\n",
       " 'attblack',\n",
       " 'wrkstat_mean',\n",
       " 'hrs1_mean',\n",
       " 'hrs2_mean',\n",
       " 'evwork_mean',\n",
       " 'occ_mean',\n",
       " 'prestige_mean',\n",
       " 'wrkslf_mean',\n",
       " 'wrkgovt_mean',\n",
       " 'commute_mean',\n",
       " 'occ80_mean',\n",
       " 'prestg80_mean',\n",
       " 'indus80_mean',\n",
       " 'marital_mean',\n",
       " 'agewed_mean',\n",
       " 'divorce_mean',\n",
       " 'widowed_mean',\n",
       " 'spwrksta_mean',\n",
       " 'sphrs1_mean',\n",
       " 'sphrs2_mean',\n",
       " 'spevwork_mean',\n",
       " 'spocc80_mean',\n",
       " 'sppres80_mean',\n",
       " 'spind80_mean',\n",
       " 'sibs_mean',\n",
       " 'childs_mean',\n",
       " 'age_mean',\n",
       " 'agekdbrn_mean',\n",
       " 'educ_mean',\n",
       " 'paeduc_mean',\n",
       " 'maeduc_mean',\n",
       " 'speduc_mean',\n",
       " 'degree_mean',\n",
       " 'padeg_mean',\n",
       " 'madeg_mean',\n",
       " 'spdeg_mean',\n",
       " 'sex_mean',\n",
       " 'race_mean',\n",
       " 'res16_mean',\n",
       " 'reg16_mean',\n",
       " 'mobile16_mean',\n",
       " 'family16_mean',\n",
       " 'mawork_mean',\n",
       " 'mawkborn_mean',\n",
       " 'born_mean',\n",
       " 'parborn_mean',\n",
       " 'granborn_mean',\n",
       " 'hompop_mean',\n",
       " 'babies_mean',\n",
       " 'preteen_mean',\n",
       " 'teens_mean',\n",
       " 'adults_mean',\n",
       " 'unrelat_mean',\n",
       " 'earnrs_mean',\n",
       " 'income_mean',\n",
       " 'rincome_mean',\n",
       " 'income86_mean',\n",
       " 'partyid_mean',\n",
       " 'polviews_mean',\n",
       " 'temp',\n",
       " '_mergescore',\n",
       " 'prestige_miss',\n",
       " 'degree_miss',\n",
       " 'childs_miss',\n",
       " 'prestg80_miss',\n",
       " 'occ_miss',\n",
       " 'marital_miss',\n",
       " 'hompop_miss',\n",
       " 'educ_miss',\n",
       " 'babies_miss',\n",
       " 'sppres80_miss',\n",
       " 'age_miss',\n",
       " 'teens_miss',\n",
       " 'preteen_miss',\n",
       " 'adults_miss',\n",
       " 'earnrs_miss',\n",
       " 'income_miss',\n",
       " 'partyid_miss',\n",
       " 'attblack_miss',\n",
       " 'd_1986',\n",
       " 'd_1987',\n",
       " 'd_1988',\n",
       " 'd_1989',\n",
       " 'd_1990',\n",
       " 'd_1991',\n",
       " 'd_1992',\n",
       " 'd_1993',\n",
       " 'd_1994',\n",
       " 'd_1995',\n",
       " 'd_1996',\n",
       " 'd_1997',\n",
       " 'd_1998',\n",
       " 'd_1999',\n",
       " 'd_2000',\n",
       " 'd_2001',\n",
       " 'd_2002',\n",
       " 'd_2003',\n",
       " 'd_2004',\n",
       " 'd_2005',\n",
       " 'd_2006',\n",
       " 'd_2007',\n",
       " 'd_2008',\n",
       " 'd_2009',\n",
       " 'd_2010']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check which variables we are working with\n",
    "list(df.columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1286cfe",
   "metadata": {},
   "source": [
    "We want to select our relevant variables, which in this case, is\n",
    "- our **outcome** variable ($Y$), (denoted by ```'y'``` in this dataset)\n",
    "- our **treatment** variable ($W$) (denoted by ```'w'``` in this dataset)\n",
    "- our relevant pre-treatment **covariates** ($X$) (in this case ```'marital', 'age', 'sex', 'educ', 'polviews', 'partyid', 'race'```)\n",
    "\n",
    "We then can subset our dataframe to only include relevant columns in our dataframe, drop observations (rows) with missing information, and cast all variables to numeric for ease of analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5919e05b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign variables\n",
    "outcome = 'y'\n",
    "treatment = 'w'\n",
    "covariates = ['marital', 'age', 'sex', 'educ', 'polviews', 'partyid', 'race']\n",
    "\n",
    "all_variable_names = [outcome, treatment] + covariates\n",
    "\n",
    "# Select the relevant columns from the dataframe and remove all others\n",
    "df = df[all_variable_names]\n",
    "\n",
    "# Remove rows with N/A values\n",
    "# In this dataframe, values that are missing/undefined are denoted as -999\n",
    "df = df[df != -999].dropna().reset_index(drop=True) \n",
    "\n",
    "# Cast all variables to numeric\n",
    "df = df.apply(pd.to_numeric, errors='coerce')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63a4e5e4",
   "metadata": {},
   "source": [
    "Now let's check out what a sample of our dataset looks like. Each row is a sample that has an outcome ($Y_i$), a treatment assignment ($W$), and covariates ($X$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e878208e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>y</th>\n",
       "      <th>w</th>\n",
       "      <th>marital</th>\n",
       "      <th>age</th>\n",
       "      <th>sex</th>\n",
       "      <th>educ</th>\n",
       "      <th>polviews</th>\n",
       "      <th>partyid</th>\n",
       "      <th>race</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>28.0</td>\n",
       "      <td>1</td>\n",
       "      <td>14.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>54.0</td>\n",
       "      <td>2</td>\n",
       "      <td>16.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>6.0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>44.0</td>\n",
       "      <td>2</td>\n",
       "      <td>16.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>77.0</td>\n",
       "      <td>2</td>\n",
       "      <td>14.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>44.0</td>\n",
       "      <td>2</td>\n",
       "      <td>14.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   y  w  marital   age  sex  educ  polviews  partyid  race\n",
       "0  0  0      5.0  28.0    1  14.0       4.0      3.0     1\n",
       "1  1  0      2.0  54.0    2  16.0       6.0      6.0     1\n",
       "2  1  0      5.0  44.0    2  16.0       2.0      0.0     1\n",
       "3  0  0      4.0  77.0    2  14.0       4.0      0.0     2\n",
       "4  0  0      4.0  44.0    2  14.0       4.0      0.0     2"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b3325a9",
   "metadata": {},
   "source": [
    "Before we get started, we might want to see the summary statistics for the data. For instance see in this case that treatment was assigned to about half (54%) of the sample rows, as seen by the mean of ```'w'``` being $.5350$.\n",
    "\n",
    "For more information about what these data mean, consult the dataset or codebook as available. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f35541cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>y</th>\n",
       "      <th>w</th>\n",
       "      <th>marital</th>\n",
       "      <th>age</th>\n",
       "      <th>sex</th>\n",
       "      <th>educ</th>\n",
       "      <th>polviews</th>\n",
       "      <th>partyid</th>\n",
       "      <th>race</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "      <td>33269.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.251916</td>\n",
       "      <td>0.535003</td>\n",
       "      <td>2.401365</td>\n",
       "      <td>46.065376</td>\n",
       "      <td>1.560702</td>\n",
       "      <td>13.117016</td>\n",
       "      <td>4.121918</td>\n",
       "      <td>2.815113</td>\n",
       "      <td>1.271274</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.434120</td>\n",
       "      <td>0.498781</td>\n",
       "      <td>1.632836</td>\n",
       "      <td>17.376747</td>\n",
       "      <td>0.496309</td>\n",
       "      <td>3.072417</td>\n",
       "      <td>1.355477</td>\n",
       "      <td>2.045163</td>\n",
       "      <td>0.571842</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>18.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>32.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>12.000000</td>\n",
       "      <td>3.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>2.000000</td>\n",
       "      <td>43.000000</td>\n",
       "      <td>2.000000</td>\n",
       "      <td>13.000000</td>\n",
       "      <td>4.000000</td>\n",
       "      <td>3.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>4.000000</td>\n",
       "      <td>59.000000</td>\n",
       "      <td>2.000000</td>\n",
       "      <td>15.000000</td>\n",
       "      <td>5.000000</td>\n",
       "      <td>5.000000</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>1.000000</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>5.000000</td>\n",
       "      <td>89.000000</td>\n",
       "      <td>2.000000</td>\n",
       "      <td>20.000000</td>\n",
       "      <td>7.000000</td>\n",
       "      <td>7.000000</td>\n",
       "      <td>3.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                  y             w       marital           age           sex  \\\n",
       "count  33269.000000  33269.000000  33269.000000  33269.000000  33269.000000   \n",
       "mean       0.251916      0.535003      2.401365     46.065376      1.560702   \n",
       "std        0.434120      0.498781      1.632836     17.376747      0.496309   \n",
       "min        0.000000      0.000000      1.000000     18.000000      1.000000   \n",
       "25%        0.000000      0.000000      1.000000     32.000000      1.000000   \n",
       "50%        0.000000      1.000000      2.000000     43.000000      2.000000   \n",
       "75%        1.000000      1.000000      4.000000     59.000000      2.000000   \n",
       "max        1.000000      1.000000      5.000000     89.000000      2.000000   \n",
       "\n",
       "               educ      polviews       partyid          race  \n",
       "count  33269.000000  33269.000000  33269.000000  33269.000000  \n",
       "mean      13.117016      4.121918      2.815113      1.271274  \n",
       "std        3.072417      1.355477      2.045163      0.571842  \n",
       "min        0.000000      1.000000      0.000000      1.000000  \n",
       "25%       12.000000      3.000000      1.000000      1.000000  \n",
       "50%       13.000000      4.000000      3.000000      1.000000  \n",
       "75%       15.000000      5.000000      5.000000      1.000000  \n",
       "max       20.000000      7.000000      7.000000      3.000000  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Calculate summary statistics\n",
    "summ_stats = pd.DataFrame(df.describe().transpose())\n",
    "summ_stats.transpose()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f99f194",
   "metadata": {},
   "source": [
    "Next, we should set up training and testing sets. We will use the training set to fit the model and the testing set to see how our model performs on data it has not seen. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ee24c37a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up a training and testing set from the scikit-learn package\n",
    "train_fraction = 0.80  # Use train_fraction % of the dataset to train our models and the rest to test\n",
    "df_train, df_test = train_test_split(df, train_size=train_fraction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06762d24",
   "metadata": {},
   "source": [
    "# 3. ***Step 1***: Estimation of Nuisance Components <a id=\"section-3\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0367b248",
   "metadata": {},
   "source": [
    "The first step is to estimate the nuisance components:\n",
    "\n",
    "- Estimating $g(x, z)$ involves fitting a model that predicts the treatment $W$, given the covariates $X$ and any additional controls $Z$. We call this the propensity estimator or ```propensity.forest``` in the code.\n",
    "<center>$g(x, z) = E[W | X = x, Z = z]$</center>\n",
    "\n",
    "\n",
    "- Estimating $m(x, w)$ involves fitting a model that predicts the potential outcome $Y$ under different levels of treatment, given the covariates $X$. We call this the outcome estimator or ```outcome.forest``` in the code.\n",
    "<center>$m(x, w) = E[Y | X = x, W = w]$</center>\n",
    "\n",
    "\n",
    "Note that the choice of estimator for $m(x, w)$ and $g(x, z)$ in this step and the choice of policy class $Π$ along with the optimizer used in Step 3 can be made fully independently."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb356bc5",
   "metadata": {},
   "source": [
    "We will be using a **Regression Forest** from R's ```grf``` package to fit a model and estimate the nuisance components. A regression forest is is an ensemble method that utilizes multiple decision trees trained on different subsets of the dataset. By employing averaging techniques, it enhances predictive accuracy and mitigates overfitting issues. For more information about ```regression_forest```, see https://www.rdocumentation.org/packages/grf/versions/0.10.2/topics/regression_forest\n",
    "\n",
    "We can use this R method in Python by using the Python library ```rpy2```. We grab our $X$, $Y$, and $X$ data from the training set, fit the models, and then use the fitted model to make $\\hat{Y}$ and $\\hat{W}$ predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ee1b731c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Activate rpy2\n",
    "pandas2ri.activate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4445c1fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grab our data for the covariate variables, outcome variable, and treatment variable from our training and test sets\n",
    "X = df_train[covariates]\n",
    "Y = df_train[outcome]\n",
    "W = df_train[treatment]\n",
    "X_test = df_test[covariates]\n",
    "Y_test = df_test[outcome]\n",
    "W_test = df_test[treatment]\n",
    "\n",
    "# Convert DataFrame to R objects\n",
    "X_r = pandas2ri.py2rpy(X)\n",
    "Y_r = pandas2ri.py2rpy(Y)\n",
    "W_r = pandas2ri.py2rpy(W)\n",
    "X_test_r = pandas2ri.py2rpy(X_test)\n",
    "\n",
    "# Assign R objects to global environment\n",
    "robjects.globalenv[\"X_r\"] = X_r\n",
    "robjects.globalenv[\"Y_r\"] = Y_r\n",
    "robjects.globalenv[\"W_r\"] = W_r\n",
    "robjects.globalenv[\"X_test_r\"] = X_test_r\n",
    "\n",
    "#Fit our regression forests in R and find W_hat and Y_hat for training and test sets\n",
    "robjects.r('''\n",
    "library(grf)\n",
    "tune.parameters = c(\"mtry\", \"min.node.size\", \"alpha\", \"imbalance.penalty\")\n",
    "propensity.forest = regression_forest(X_r, W_r, tune.parameters = tune.parameters)\n",
    "W_hat = predict(propensity.forest)$predictions\n",
    "W_hat.test <- predict(propensity.forest, newdata=X_test_r)$predictions\n",
    "\n",
    "outcome.forest = regression_forest(X_r, Y_r, tune.parameters = tune.parameters)\n",
    "Y_hat = predict(outcome.forest)$predictions\n",
    "Y_hat.test = predict(outcome.forest, newdata=X_test_r)$predictions\n",
    "''')\n",
    "\n",
    "# Convert Y_hat and W_hat to Python objects\n",
    "Y_hat = robjects.r['Y_hat'].flatten()\n",
    "W_hat = robjects.r['W_hat'].flatten()\n",
    "Y_hat_test = robjects.r['Y_hat.test'].flatten()\n",
    "W_hat_test = robjects.r['W_hat.test'].flatten()\n",
    "\n",
    "# Convert Y_hat and W_hat to NumPy Arrays for ease of use in Python\n",
    "Y_hat = np.array(Y_hat)\n",
    "W_hat = np.array(W_hat)\n",
    "Y_hat_test = np.array(Y_hat_test)\n",
    "W_hat_test = np.array(W_hat_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffea949f",
   "metadata": {},
   "source": [
    "# 4. ***Step 2***: Computing Doubly Robust Scores <a id=\"section-4\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aa651fa",
   "metadata": {},
   "source": [
    "We have now fit models and estimated our nuisance paramenters. The next step is to compute doubly robust scores."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f770294",
   "metadata": {},
   "source": [
    "To predict treatment effects, we can use the **Causal Forest** model, which is implemented in ```grf```'s ```causal_forest``` in R, transferred to Python using ```rpy2```. This model builds upon the random forest framework by training multiple decision trees on different subsets of the data. In addition to standard predictions, causal forests specifically focus on estimating the causal effect of a treatment or intervention on an outcome variable. They leverage the heterogeneity present in the data and incorporate the propensity score, which represents the probability of treatment assignment, to estimate average treatment effects and other causal effects of interest. You can find more information here: https://grf-labs.github.io/grf/reference/causal_forest.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba699e96",
   "metadata": {},
   "source": [
    "We use these treatment effects to find our doubly robust scores, $\\hat{\\Gamma}_i$.\n",
    "\n",
    "<center>$\\hat{\\Gamma}_i = \\hat{\\tau}_{-i}(X_i) + W_i - \\hat{e}_{-i}(X_i)\\hat{e}_{-i}(X_i)(1 - \\hat{e}_{-i}(X_i))(Y_i - \\hat{\\mu}_{-i}W_i)$</center>\n",
    "\n",
    "In the above equation, \n",
    "- $\\hat{\\tau}(\\cdot)$ treatment effect estimates, \n",
    "- $\\hat{e}(\\cdot)$: propensity score estimates\n",
    "- $\\hat{\\mu}_0$, $\\hat{\\mu}_1$: estimates of $E[Y|X_i,W_i=0]$ and $E[Y|X_i,W_i=1]$ respectively. \n",
    "- *Note:* $-i$ indicates that the estimate for observation $i$ is an out-of-bag (OOB) estimate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd12f70f",
   "metadata": {},
   "source": [
    "The use of doubly robust scores offers two **key advantages**:\n",
    "\n",
    "1. They provide protection against model misspecification: if either the outcome model or the propensity score model is correctly specified, the average of the doubly robust scores consistently estimates the average treatment effect. \n",
    "\n",
    "2. Their average achieves the semi-parametric lower bound as an estimation of the average treatment effect. This semi-parametric efficiency translates into provable theoretical guarantees on the performance of the optimal estimated policy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa21ae1",
   "metadata": {},
   "source": [
    "We will compute scores for the training and testing sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ae8a9b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In R, fit a causal forest model and obtain estimated treatment effects\n",
    "robjects.r(\n",
    "    '''\n",
    "    library(grf)\n",
    "    tune.parameters = c(\"mtry\", \"min.node.size\", \"alpha\", \"imbalance.penalty\")\n",
    "    cf = causal_forest(X_r, Y_r, W_r, Y.hat = Y_hat, W.hat = W_hat, tune.parameters = tune.parameters)\n",
    "\n",
    "    # Obtain the estimated treatment effect\n",
    "    tauhat = predict(cf)$predictions\n",
    "    tauhat.test <- predict(cf, newdata=X_test_r)$predictions\n",
    "    ''')\n",
    "\n",
    "# Convert output to Python NumPy array\n",
    "tau_hat = robjects.r['tauhat'].flatten()\n",
    "tau_hat = np.array(tau_hat)\n",
    "tau_hat_test = robjects.r['tauhat.test'].flatten()\n",
    "tau_hat_test = np.array(tau_hat_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4fe680b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate E[Y|X=x, W=0] and E[Y|X=x, W=1] for training set\n",
    "mu_hat_0 = Y_hat - W_hat * tau_hat\n",
    "mu_hat_1 = Y_hat + (1 - W_hat) * tau_hat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7ee1d33e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Estimate Gamma_hat using our mu_hat_0 and mu_hat_1 for training set\n",
    "resid = Y - W * mu_hat_1 - (1 - W) * mu_hat_0\n",
    "weights = (W - W_hat) / (W_hat * (1 - W_hat))\n",
    "Gamma_hat = tau_hat + weights * resid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c07022a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11540   -0.446005\n",
       "19699   -0.348049\n",
       "32634   -0.503020\n",
       "30150    0.537836\n",
       "17196   -0.469412\n",
       "           ...   \n",
       "8663    -0.575906\n",
       "27869   -0.705942\n",
       "2642    -0.472470\n",
       "29070   -0.451484\n",
       "30814   -0.395903\n",
       "Length: 26615, dtype: float64"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Gamma_hat #Let's look at our scores for the training set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cd9885fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate E[Y|X=x, W=0] and E[Y|X=x, W=1] for test set\n",
    "mu_hat_0_test = Y_hat_test - W_hat_test * tau_hat_test\n",
    "mu_hat_1_test = Y_hat_test + (1 - W_hat_test) * tau_hat_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0be31d77",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Estimate Gamma_hat using our mu_hat_0 and mu_hat_1 for test set\n",
    "resid_test = Y_test - W_test * mu_hat_1_test - (1 - W_test) * mu_hat_0_test\n",
    "weights_test = (W_test - W_hat_test) / (W_hat_test * (1 - W_hat_test))\n",
    "Gamma_hat_test = tau_hat_test + weights_test * resid_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "744aba92",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "29590   -0.298898\n",
       "17002   -0.201816\n",
       "18485    0.757932\n",
       "17866    0.995479\n",
       "5410    -1.468764\n",
       "           ...   \n",
       "21142   -0.390404\n",
       "8342    -0.425066\n",
       "6709    -0.766035\n",
       "30450    0.550738\n",
       "11857    1.331888\n",
       "Length: 6654, dtype: float64"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Gamma_hat_test #Let's look at our test set scores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d78ede45",
   "metadata": {},
   "source": [
    "Now that we have obtained the scores, we subtract the treatment cost. In real world applications, the true cost of treatment would be subtracted. However, for illustrative purposes, we calibrate the cost for this tutorial by setting it equal to the median of the CATE (Conditional Average Treatment Effect) estimates. This calibration ensures that a portion of the population is assigned to receive treatment while others are not, thereby creating a nontrivial policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3f55b275",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define cost to be median CATE\n",
    "cost = np.median(tau_hat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8ccb270e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Subtract the cost from the Gamma scores\n",
    "Gamma_hat_net = Gamma_hat - cost\n",
    "Gamma_hat_test_net = Gamma_hat_test - cost"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6de91584",
   "metadata": {},
   "source": [
    "# 5. ***Step 3*** Fitting the Policy Tree <a id=\"section-5\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d5156f0",
   "metadata": {},
   "source": [
    "Next, we prepare our data for the policy tree. A policy tree is a decision tree-based algorithm that is specifically designed to learn optimal decision-making policies in sequential decision problems. Given doubly robust reward estimates, we use policy trees to find a rule-based treatment assignment policy, where the policy takes the form of a shallow decision tree that is globally (or close to) optimal. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cd4ffef",
   "metadata": {},
   "source": [
    "Using doubly robust scores we found, we now aim to find a policy in the set of feasible policies ($\\pi \\in \\Pi$) that maximizes the objective function:\n",
    "\n",
    "<center>$\\pi^* = \\arg\\max_{\\pi \\in \\Pi} \\left[ \\frac{1}{n} \\sum_{i=1}^{n} \\Gamma_i(\\pi(X_i)) \\right]$</center>\n",
    "where:\n",
    "\n",
    "- $\\pi^*$ represents the estimated optimal policy,\n",
    "- $\\arg\\max_{\\pi \\in \\Pi}$ denotes the maximization over the set of policies $\\Pi$,\n",
    "- $\\frac{1}{n} \\sum_{i=1}^{n} \\Gamma_i(\\pi(X_i))$ is the average of the doubly robust scores $\\Gamma_i(\\pi(X_i))$ over all observations.\n",
    "\n",
    "\n",
    "In this tutorial, we will be using the R package ```policytree```, developed by Athey and Wager and co-authors and used in \"Policy Learning with Observational Data\" (2021) for this precise purpose. For more information, see https://cran.r-project.org/web/packages/policytree/index.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ada4812",
   "metadata": {},
   "source": [
    "We should determine what our class of feasible policies is. Let's first look at a policy tree that is constrained to have a depth of 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c21cf638",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"184pt\" height=\"122pt\" viewBox=\"0.00 0.00 184.11 121.60\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 117.6)\">\n",
       "<title>nodes</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"transparent\" points=\"-4,4 -4,-117.6 180.105,-117.6 180.105,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"133.9992,-113.6 41.1058,-113.6 41.1058,-77.6 133.9992,-77.6 133.9992,-113.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"87.5525\" y=\"-91.4\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> partyid &lt;= 3 </text>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>1</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"79.1576,-41.4019 -.0526,-41.4019 -.0526,-.1981 79.1576,-.1981 79.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"39.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"39.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 2 </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;1 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>0-&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M75.933,-77.493C70.6766,-69.3016 64.329,-59.41 58.4174,-50.1978\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"61.2796,-48.1774 52.9331,-41.6515 55.3883,-51.958 61.2796,-48.1774\"/>\n",
       "<text text-anchor=\"middle\" x=\"47.6026\" y=\"-61.8766\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">True</text>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>2</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"176.1576,-41.4019 96.9474,-41.4019 96.9474,-.1981 176.1576,-.1981 176.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"136.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"136.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 1 </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;2 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>0-&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M99.414,-77.493C104.78,-69.3016 111.2599,-59.41 117.2946,-50.1978\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"120.3411,-51.9344 122.8931,-41.6515 114.4856,-48.0986 120.3411,-51.9344\"/>\n",
       "<text text-anchor=\"middle\" x=\"127.9935\" y=\"-61.9257\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">False</text>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Set up our Gamma hat nets for implementation in the policy tree\n",
    "Gamma_hat_net_r = pandas2ri.py2rpy(Gamma_hat_net)\n",
    "robjects.globalenv[\"Gamma_hat_net_r\"] = Gamma_hat_net_r\n",
    "\n",
    "#Use policytree to find optimal tree that is within class (depth 1)\n",
    "robjects.r(\n",
    "    '''\n",
    "    library(policytree)\n",
    "    library(DiagrammeRsvg)\n",
    "    \n",
    "    # policy_tree takes in our X, our doubly robust estimates, and our chosen tree depth\n",
    "    tree.1 <- policy_tree(X_r, cbind(-Gamma_hat_net_r, Gamma_hat_net_r), depth = 1)\n",
    "    policy.1 <- predict(tree.1, X_r) - 1\n",
    "    tree.plot = plot(tree.1, leaf.names=c('control', 'treatment'))\n",
    "    cat(DiagrammeRsvg::export_svg(tree.plot), file = 'plot.svg')\n",
    "    ''')\n",
    "\n",
    "# Read the SVG file\n",
    "with open('plot.svg', 'r') as svg_file:\n",
    "    svg_data = svg_file.read()\n",
    "\n",
    "# Display the SVG\n",
    "display(SVG(svg_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fc8fce9",
   "metadata": {},
   "source": [
    "As we see in this plot, we get an optimal policy tree that splits on the variable partyid, with units with partyid <= 3 receiving treatment and units with partyid > 3 recieving control.\n",
    "\n",
    "- **```Action 1``` denotes control and ```Action 2``` denotes treatment.**\n",
    "- If you would like to see the predicted assignments for each row, you can print ```policy.1```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2152ec5",
   "metadata": {},
   "source": [
    "Now let's say we want to constrain our feasible policies to depth of 2 trees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e6803a27",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"380pt\" height=\"194pt\" viewBox=\"0.00 0.00 380.11 193.60\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 189.6)\">\n",
       "<title>nodes</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"transparent\" points=\"-4,4 -4,-189.6 376.105,-189.6 376.105,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"228.7712,-185.6 142.3338,-185.6 142.3338,-149.6 228.7712,-149.6 228.7712,-185.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"185.5525\" y=\"-163.4\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> educ &lt;= 13 </text>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>1</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"172.6462,-113.6 96.4588,-113.6 96.4588,-77.6 172.6462,-77.6 172.6462,-113.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"134.5525\" y=\"-91.4\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> race &lt;= 1 </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;1 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>0-&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M172.683,-149.4314C166.869,-141.2232 159.8829,-131.3606 153.4879,-122.3323\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"156.2315,-120.1505 147.5952,-114.0133 150.5194,-124.1966 156.2315,-120.1505\"/>\n",
       "<text text-anchor=\"middle\" x=\"143.3878\" y=\"-134.4567\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">True</text>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>2</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"282.9992,-113.6 190.1058,-113.6 190.1058,-77.6 282.9992,-77.6 282.9992,-113.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"236.5525\" y=\"-91.4\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> partyid &lt;= 3 </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;2 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>0-&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M198.422,-149.4314C204.236,-141.2232 211.2221,-131.3606 217.6171,-122.3323\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"220.5856,-124.1966 223.5098,-114.0133 214.8735,-120.1505 220.5856,-124.1966\"/>\n",
       "<text text-anchor=\"middle\" x=\"227.7172\" y=\"-134.4567\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">False</text>\n",
       "</g>\n",
       "<!-- 3 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>3</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"79.1576,-41.4019 -.0526,-41.4019 -.0526,-.1981 79.1576,-.1981 79.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"39.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"39.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 1 </text>\n",
       "</g>\n",
       "<!-- 1&#45;&gt;3 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>1-&gt;3</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M111.5556,-77.493C100.2942,-68.6261 86.5023,-57.7668 74.0138,-47.9338\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"76.0571,-45.0879 66.035,-41.6515 71.7267,-50.5877 76.0571,-45.0879\"/>\n",
       "</g>\n",
       "<!-- 4 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"176.1576,-41.4019 96.9474,-41.4019 96.9474,-.1981 176.1576,-.1981 176.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"136.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"136.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 2 </text>\n",
       "</g>\n",
       "<!-- 1&#45;&gt;4 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>1-&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M135.0366,-77.493C135.2421,-69.8083 135.4876,-60.6272 135.7206,-51.9145\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"139.2263,-51.7415 135.995,-41.6515 132.2288,-51.5543 139.2263,-51.7415\"/>\n",
       "</g>\n",
       "<!-- 5 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>5</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"275.1576,-41.4019 195.9474,-41.4019 195.9474,-.1981 275.1576,-.1981 275.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"235.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"235.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 2 </text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;5 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>2-&gt;5</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M236.3104,-77.493C236.2077,-69.8083 236.0849,-60.6272 235.9685,-51.9145\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"239.4647,-51.6038 235.8313,-41.6515 232.4654,-51.6974 239.4647,-51.6038\"/>\n",
       "</g>\n",
       "<!-- 6 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>6</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"372.1576,-41.4019 292.9474,-41.4019 292.9474,-.1981 372.1576,-.1981 372.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"332.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"332.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 1 </text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;6 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>2-&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M259.7914,-77.493C271.1714,-68.6261 285.1085,-57.7668 297.7284,-47.9338\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"300.0542,-50.5587 305.7912,-41.6515 295.7518,-45.0369 300.0542,-50.5587\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#Use policytree to find optimal tree that is within class (depth 2)\n",
    "robjects.r(\n",
    "    '''\n",
    "    tree.2 <- policy_tree(X_r, cbind(-Gamma_hat_net_r, Gamma_hat_net_r), depth = 2)\n",
    "    policy.2 <- predict(tree.2, X_r) - 1\n",
    "    tree.2.plot = plot(tree.2, leaf.names=c('control', 'treatment'))\n",
    "    cat(DiagrammeRsvg::export_svg(tree.2.plot), file = 'plot2.svg')\n",
    "    ''')\n",
    "\n",
    "# Read the SVG file\n",
    "with open('plot2.svg', 'r') as svg_file:\n",
    "    svg_data = svg_file.read()\n",
    "\n",
    "# Display the SVG\n",
    "display(SVG(svg_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "637bdcea",
   "metadata": {},
   "source": [
    "In this case, we first split on education. \n",
    "- If education is <= 13 and race <= 1, we assign control. \n",
    "- If education is <= 13 and race > 1, we assign treatment. \n",
    "- If education is > 14 and partyid <= 3, we assign treatment.\n",
    "- If education is > 14 and partyid > 3, we assign control. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6176369",
   "metadata": {},
   "source": [
    "In this case, it looks like we are splitting on race. Perhaps for ethical or legal reasons, we don't want our tree to split on race, age, or gender variables. We can limit our splitting variables by making a subset of the dataframe with possible split variables and performing the analysis on them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7a9ba33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sensitive_features = [\"sex\", \"race\", \"age\"]\n",
    "safe_features = [feature for feature in X.columns if feature not in sensitive_features]\n",
    "X_safe = X[safe_features]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "62ef819e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['marital', 'educ', 'polviews', 'partyid'], dtype='object')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_safe.columns # Our possible splitting variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "13f47f17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<svg xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\" width=\"384pt\" height=\"194pt\" viewBox=\"0.00 0.00 384.11 193.60\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 189.6)\">\n",
       "<title>nodes</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"transparent\" points=\"-4,4 -4,-189.6 380.105,-189.6 380.105,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"230.7712,-185.6 144.3338,-185.6 144.3338,-149.6 230.7712,-149.6 230.7712,-185.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"187.5525\" y=\"-163.4\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> educ &lt;= 12 </text>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>1</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"178.995,-113.6 86.11,-113.6 86.11,-77.6 178.995,-77.6 178.995,-113.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"132.5525\" y=\"-91.4\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> marital &lt;= 1 </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;1 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>0-&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M173.6737,-149.4314C167.3389,-141.1386 159.7141,-131.157 152.7599,-122.0533\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"155.47,-119.8353 146.6182,-114.0133 149.9073,-124.0847 155.47,-119.8353\"/>\n",
       "<text text-anchor=\"middle\" x=\"143.3013\" y=\"-134.5923\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">True</text>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>2</title>\n",
       "<polygon fill=\"none\" stroke=\"#000000\" points=\"288.9992,-113.6 196.1058,-113.6 196.1058,-77.6 288.9992,-77.6 288.9992,-113.6\"/>\n",
       "<text text-anchor=\"middle\" x=\"242.5525\" y=\"-91.4\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> partyid &lt;= 3 </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;2 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>0-&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M201.4313,-149.4314C207.7661,-141.1386 215.3909,-131.157 222.3451,-122.0533\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"225.1977,-124.0847 228.4868,-114.0133 219.635,-119.8353 225.1977,-124.0847\"/>\n",
       "<text text-anchor=\"middle\" x=\"231.8037\" y=\"-134.5923\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">False</text>\n",
       "</g>\n",
       "<!-- 3 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>3</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"79.1576,-41.4019 -.0526,-41.4019 -.0526,-.1981 79.1576,-.1981 79.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"39.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"39.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 1 </text>\n",
       "</g>\n",
       "<!-- 1&#45;&gt;3 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>1-&gt;3</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M110.0398,-77.493C99.0154,-68.6261 85.5139,-57.7668 73.2883,-47.9338\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"75.4634,-45.1916 65.4775,-41.6515 71.0762,-50.6462 75.4634,-45.1916\"/>\n",
       "</g>\n",
       "<!-- 4 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"176.1576,-41.4019 96.9474,-41.4019 96.9474,-.1981 176.1576,-.1981 176.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"136.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"136.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 2 </text>\n",
       "</g>\n",
       "<!-- 1&#45;&gt;4 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>1-&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M133.5208,-77.493C133.9317,-69.8083 134.4227,-60.6272 134.8886,-51.9145\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"138.3984,-51.8242 135.4374,-41.6515 131.4084,-51.4503 138.3984,-51.8242\"/>\n",
       "</g>\n",
       "<!-- 5 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>5</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"279.1576,-41.4019 199.9474,-41.4019 199.9474,-.1981 279.1576,-.1981 279.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"239.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"239.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 2 </text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;5 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>2-&gt;5</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M241.8263,-77.493C241.5181,-69.8083 241.1498,-60.6272 240.8004,-51.9145\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"244.2868,-51.5032 240.3888,-41.6515 237.2924,-51.7838 244.2868,-51.5032\"/>\n",
       "</g>\n",
       "<!-- 6 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>6</title>\n",
       "<polygon fill=\"#c1b2ff\" stroke=\"#c1b2ff\" points=\"376.1576,-41.4019 296.9474,-41.4019 296.9474,-.1981 376.1576,-.1981 376.1576,-41.4019\"/>\n",
       "<text text-anchor=\"middle\" x=\"336.5525\" y=\"-25\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> leaf node</text>\n",
       "<text text-anchor=\"middle\" x=\"336.5525\" y=\"-8.2\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\"> action = 1 </text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;6 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>2-&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M265.3073,-77.493C276.4502,-68.6261 290.0969,-57.7668 302.4539,-47.9338\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"304.7031,-50.6169 310.3487,-41.6515 300.3445,-45.1394 304.7031,-50.6169\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>"
      ],
      "text/plain": [
       "<IPython.core.display.SVG object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Convert to R object and assign to global environment\n",
    "X_safe_r = pandas2ri.py2rpy(X_safe)\n",
    "robjects.globalenv[\"X_safe_r\"] = X_safe_r\n",
    "\n",
    "#Use policytree to find optimal tree that is within class (safe variables only; depth 2)\n",
    "robjects.r(\n",
    "    '''\n",
    "    tree.safe <- policy_tree(X_safe_r, cbind(-Gamma_hat_net_r, Gamma_hat_net_r), depth = 2)\n",
    "    policy.safe <- predict(tree.safe, X_safe_r)\n",
    "    tree.safe.plot = plot(tree.safe, leaf.names=c('control', 'treatment'))\n",
    "    cat(DiagrammeRsvg::export_svg(tree.safe.plot), file = 'plot3.svg')\n",
    "    ''')\n",
    "\n",
    "# Read the SVG file\n",
    "with open('plot3.svg', 'r') as svg_file:\n",
    "    svg_data = svg_file.read()\n",
    "\n",
    "# Display the SVG\n",
    "display(SVG(svg_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cbede78",
   "metadata": {},
   "source": [
    "Notice that we no longer are splitting on sensitive characteristics."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25e72316",
   "metadata": {},
   "source": [
    "# 6. Evaluation <a id='section-6'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "835d190d",
   "metadata": {},
   "source": [
    "Now that we have a policy, how can we evaluate it?\n",
    "\n",
    "The **improvement** of a policy $π$ over a random policy can be calculated as the difference between their values. The **value** of a policy is the expected benefit it provides over treating no one. We can express this improvement formally as:\n",
    "\n",
    "<center>$A(π) = V(π) - V(π_{random})$</center>\n",
    "\n",
    "We also know that\n",
    "- $V(π)$, which equals $E[π(X_i)τ(X_i)]$, is the value of policy π, which is the expected treatment effect when applying the policy\n",
    "- $V(π_{random})$, which equals $\\frac{1}{2}E[τ(Xi)]+\\frac{1}{2}0 = \\frac{1}{2} E[τ(Xi)]$, is the value when treatment and non-treatment are assigned with equal probability.\n",
    "\n",
    "To calculate the improvement, we substitute the expressions for the values into the equation (scaled by 2 for simplification convenience):\n",
    "\n",
    "$A(\\pi) = 2(V(\\pi) - V(\\pi_{\\text{random}})) $\n",
    "\n",
    "$= 2(E[\\pi(X_i)\\tau(X_i)] - \\frac{1}{2}E[\\tau(X_i)]) $\n",
    "\n",
    "$= 2E[\\tau(X_i)|\\pi(X_i)=1]P(\\pi(X_i)=1) - (E[\\tau(X_i)|\\pi(X_i)=1]P(\\pi(X_i)=1) + E[\\tau(X_i)|\\pi(X_i)=0]P(\\pi(X_i)=0) $\n",
    "\n",
    "$= E[\\tau(X_i)|\\pi(X_i)=1]P(\\pi(X_i)=1) - E[\\tau(X_i)|\\pi(X_i)=0]P(\\pi(X_i)=0)$\n",
    "\n",
    "\n",
    "Each term in the expression represents the conditional expectation of the treatment effect under the corresponding policy assignment $(π(X_i)=1$ or $π(X_i)=0)$ weighted by the probability of that assignment. This simplified expression represents the difference in the expected treatment effects for those assigned treatment minus the average treatment effect for those assigned no treatment, weighted by their respective probabilities. We will use this method to evaluate our model.\n",
    "\n",
    "We will use ```policytree```'s ```.predict``` method to predict optimal assignment and then calculate the value of the policy over random."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "184778e8",
   "metadata": {},
   "source": [
    "For illustration purposes, we will show evaluation for the tree with ```depth = 2```. We use the ```psych``` R packages to give us descriptive statistics about evaluation. We will output the doubly robust estimated benefit of optimal tree policy versus random policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2f7f539f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                        A.hat         SE   Lower CI   Upper CI\n",
      "Optimal Tree Performance - Train  0.003500206 0.02610541 -0.0476664 0.05466681\n",
      "Optimal Tree Performance - Test  -0.004116695 0.05241322 -0.1068466 0.09861322\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>A.hat</th>\n",
       "      <th>SE</th>\n",
       "      <th>Lower CI</th>\n",
       "      <th>Upper CI</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Optimal Tree Performance - Train</th>\n",
       "      <td>0.003500</td>\n",
       "      <td>0.026105</td>\n",
       "      <td>-0.047666</td>\n",
       "      <td>0.054667</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Optimal Tree Performance - Test</th>\n",
       "      <td>-0.004117</td>\n",
       "      <td>0.052413</td>\n",
       "      <td>-0.106847</td>\n",
       "      <td>0.098613</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                     A.hat        SE  Lower CI  Upper CI\n",
       "Optimal Tree Performance - Train  0.003500  0.026105 -0.047666  0.054667\n",
       "Optimal Tree Performance - Test  -0.004117  0.052413 -0.106847  0.098613"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Convert to R object and assign to global environment\n",
    "Gamma_hat_net_r = pandas2ri.py2rpy(Gamma_hat_net)\n",
    "robjects.globalenv[\"Gamma_hat_net_r\"] = Gamma_hat_net\n",
    "Gamma_hat_test_net_r = pandas2ri.py2rpy(Gamma_hat_test_net)\n",
    "robjects.globalenv[\"Gamma_hat_test_net_r\"] = Gamma_hat_test_net\n",
    "\n",
    "#Calculate Perfomances\n",
    "robjects.r(\"\"\"\n",
    "library(psych) \n",
    "tree.assignment.train <- as.numeric(as.character(predict(tree.2, newdata = X_r, type=\"node.id\")))\n",
    "tree.assignment.test <- as.numeric(as.character(predict(tree.2, newdata = X_test_r, type=\"node.id\")))\n",
    "\n",
    "# Calculate value over random policy and print out results\n",
    "tree.A.hat <- data.frame(rbind(\n",
    "  \"Optimal Tree Performance - Train\" = describe(tree.assignment.train*Gamma_hat_net_r)[c(\"mean\", \"se\")],\n",
    "  \"Optimal Tree Performance - Test\"  = describe(tree.assignment.test*Gamma_hat_test_net_r)[c(\"mean\", \"se\")]))\n",
    "tree.A.hat$lower.ci <- with(tree.A.hat, mean - 1.96 * se)\n",
    "tree.A.hat$upper.ci <- with(tree.A.hat, mean + 1.96 * se)\n",
    "colnames(tree.A.hat) <- c(\"A.hat\", \"SE\", \"Lower CI\", \"Upper CI\")\n",
    "print(tree.A.hat)\n",
    "\"\"\"         )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6efaece1",
   "metadata": {},
   "source": [
    "As a note, this policy tree was fit on the training data, so the training sample estimate may be biased upwards.\n",
    "\n",
    "We can evaluate any of our trees using this method to uncover our estimated improvement of our optimal policy trees in the policy class over random."
   ]
  }
 ],
 "metadata": {
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}