sepsis_investigation_run_estimators.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import itertools\n",
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "from stable_baselines3 import DQN, PPO\n",
    "\n",
    "from gym_sepsis.envs.sepsis_env import SepsisEnv\n",
    "from environments.sepsis_env_wrapper import SepsisEnvWrapper\n",
    "from policies.sb3_policy import SB3Policy\n",
    "from utils.offline_dataset import OfflineRLDataset\n",
    "from models.large_a_fnn_nuisance_model import LargeAFeedForwardNuisanceModel\n",
    "from models.fnn_critic import FeedForwardCritic\n",
    "from learners.robust_fqi_learner import RobustFQILearner\n",
    "from learners.iterative_sieve_critic import IterativeSieveLearner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('configs/sepsis_config.json') as f:\n",
    "    sepsis_config = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_evaluation_run(config, rep_i, adversarial_lambda, device=None):\n",
    "    print(f'doing evaluation for lambda={adversarial_lambda}, rep={rep_i}')\n",
    "\n",
    "    base_env = SepsisEnv()\n",
    "    env = SepsisEnvWrapper(base_env=base_env, s_init_idx=0)\n",
    "\n",
    "    s_dim = env.get_s_dim()\n",
    "    num_a = env.get_num_a()\n",
    "    gamma = config['gamma']\n",
    "    default_batch_size = config['default_batch_size']\n",
    "\n",
    "    pi_e = SB3Policy(env, model=DQN.load(config['dqn_model_path']))\n",
    "    pi_e_name = config['pi_e_name']\n",
    "\n",
    "    s_init, a_init = env.get_s_a_init(pi_e)\n",
    "    if device is not None:\n",
    "        s_init = s_init.to(device)\n",
    "        a_init = a_init.to(device)\n",
    "\n",
    "    model_config = config['model_config']\n",
    "    model = LargeAFeedForwardNuisanceModel(\n",
    "        s_dim=s_dim,\n",
    "        num_a=num_a,\n",
    "        gamma=gamma,\n",
    "        config=model_config\n",
    "    )\n",
    "    critic_class = FeedForwardCritic\n",
    "    critic_config = config['critic_config']\n",
    "    critic_kwargs = {\n",
    "        's_dim': s_dim,\n",
    "        'num_a': num_a,\n",
    "        'config': critic_config\n",
    "    }\n",
    "\n",
    "    train_dataset = OfflineRLDataset.load_dataset(config['train_dataset_path'])\n",
    "    test_dataset = OfflineRLDataset.load_dataset(config['test_dataset_path'])\n",
    "    if device is not None:\n",
    "        train_dataset.to(device)\n",
    "        test_dataset.to(device)\n",
    "\n",
    "    # first train q/beta\n",
    "    print('  -- training q')\n",
    "    q_learner = RobustFQILearner(\n",
    "        nuisance_model=model, gamma=gamma, use_dual_cvar=True,\n",
    "        adversarial_lambda=adversarial_lambda,\n",
    "    )\n",
    "    q_learner_kwargs = config['q_learner_kwargs']\n",
    "    q_learner.train(\n",
    "        dataset=train_dataset, pi_e_name=pi_e_name, verbose=False,\n",
    "        device=device, **q_learner_kwargs,\n",
    "    )\n",
    "    model.freeze_embeds()\n",
    "\n",
    "    # second train eta\n",
    "    print('  -- training eta')\n",
    "    eta_learner = IterativeSieveLearner(\n",
    "        nuisance_model=model, gamma=gamma, use_dual_cvar=True,\n",
    "        adversarial_lambda=adversarial_lambda,\n",
    "        train_q_beta=False, train_eta=True, train_w=False, debug_beta=False,\n",
    "    )\n",
    "    eta_learner_kwargs = config['eta_learner_kwargs']\n",
    "    eta_learner.train(\n",
    "        dataset=train_dataset, pi_e_name=pi_e_name, verbose=False,\n",
    "        device=device, init_basis_func=env.bias_basis_func,\n",
    "        num_init_basis=1, critic_class=critic_class, s_init=s_init,\n",
    "        critic_kwargs=critic_kwargs, **eta_learner_kwargs,\n",
    "    )\n",
    "\n",
    "    # third train w\n",
    "    print('  -- training w')\n",
    "    w_learner = IterativeSieveLearner(\n",
    "        nuisance_model=model, gamma=gamma, use_dual_cvar=True,\n",
    "        adversarial_lambda=adversarial_lambda,\n",
    "        train_q_beta=False, train_eta=False, train_w=True, debug_beta=False,\n",
    "    )\n",
    "    w_learner_kwargs = config['w_learner_kwargs']\n",
    "    w_learner.train(\n",
    "        dataset=train_dataset, pi_e_name=pi_e_name, verbose=False,\n",
    "        device=device, init_basis_func=env.bias_basis_func,\n",
    "        num_init_basis=1, critic_class=critic_class, s_init=s_init,\n",
    "        critic_kwargs=critic_kwargs, **w_learner_kwargs,\n",
    "    )\n",
    "\n",
    "    model_path_base = config['base_model_path']\n",
    "    model_name = 'sepsis_model'\n",
    "    model_name += f'_lambda={adversarial_lambda}'\n",
    "    model_name += f'_rep={rep_i}'\n",
    "    model_path = os.path.join(model_path_base, model_name)\n",
    "    model.save_model(model_path)\n",
    "\n",
    "    ## evaluate model using 3 policy value estimators\n",
    "\n",
    "    dl_test = test_dataset.get_batch_loader(batch_size=default_batch_size)\n",
    "\n",
    "    q_pv = model.estimate_policy_val_q(\n",
    "        s_init=s_init, a_init=a_init, gamma=gamma\n",
    "    )\n",
    "    w_pv = model.estimate_policy_val_w(\n",
    "        dl=dl_test, pi_e_name=pi_e_name,\n",
    "    )\n",
    "    w_pv_norm = model.estimate_policy_val_w(\n",
    "        dl=dl_test, pi_e_name=pi_e_name, normalize=True,\n",
    "    )\n",
    "    dr_pv = model.estimate_policy_val_dr(\n",
    "        s_init=s_init, a_init=a_init, pi_e_name=pi_e_name, dl=dl_test,\n",
    "        adversarial_lambda=adversarial_lambda, gamma=gamma, dual_cvar=True,\n",
    "    )\n",
    "    dr_pv_norm = model.estimate_policy_val_dr(\n",
    "        s_init=s_init, a_init=a_init, pi_e_name=pi_e_name, dl=dl_test,\n",
    "        adversarial_lambda=adversarial_lambda, gamma=gamma, dual_cvar=True,\n",
    "        normalize=True,\n",
    "    )\n",
    "    pv_results = {\n",
    "        'q': q_pv, 'w': w_pv, 'w_norm': w_pv_norm,\n",
    "        'dr': dr_pv, 'dr_norm': dr_pv_norm, \n",
    "    }\n",
    "    results = []\n",
    "    for key, val in pv_results.items():\n",
    "        row = {\n",
    "            'rep_i': rep_i,\n",
    "            'lambda': adversarial_lambda,\n",
    "            'est_policy_value': val,\n",
    "            'estimator': key,\n",
    "        }\n",
    "        results.append(row)\n",
    "    return results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# iterate over values of Lambda and repetition index, and run experiment\n",
    "restart_range = sepsis_config['num_restart_range']\n",
    "lambda_range = sepsis_config['adversarial_lambda_values']\n",
    "results_list = []\n",
    "\n",
    "for rep_i, adversarial_lambda in itertools.product(restart_range, lambda_range):\n",
    "    next_results = single_evaluation_run(\n",
    "        config=sepsis_config,\n",
    "        rep_i=rep_i,\n",
    "        adversarial_lambda=adversarial_lambda,\n",
    "        device=sepsis_config['device'],\n",
    "    )\n",
    "    results_list.extend(next_results)\n",
    "    results_df = pd.DataFrame(results_list)\n",
    "    display(results_df)\n",
    "    results_df.to_csv('sepsis_results.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now, build plot of results\n",
    "\n",
    "for q in (0.2, 0.5, 0.8):\n",
    "\n",
    "    keep_rows = (results_df['est_policy_value'] > 0) \\\n",
    "                & (results_df['est_policy_value'] < 1)\n",
    "    plot_df = results_df[keep_rows]\\\n",
    "                    .set_index(['rep_i', 'lambda', 'estimator'])\\\n",
    "                    .groupby(['lambda', 'estimator'])\\\n",
    "                    .quantile(q)\\\n",
    "                    .reset_index()\\\n",
    "                    .pivot(index='lambda', values='est_policy_value', columns='estimator')\\\n",
    "                    .loc[:, ['q', 'w', 'dr']]\n",
    "\n",
    "    print(f'{q} quantile')\n",
    "    display(plot_df)\n",
    "    print('')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "adversarial-rl",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}