svg_kernel.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import random\n",
    "import time\n",
    "from distutils.util import strtobool\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.distributions.normal import Normal\n",
    "from torch.distributions.multivariate_normal import MultivariateNormal\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "PERIOD = 7\n",
    "\n",
    "# Env\n",
    "import gym, json\n",
    "from gym import spaces\n",
    "from epipolicy.core.epidemic import construct_epidemic\n",
    "from epipolicy.obj.act import construct_act\n",
    "\n",
    "class EpiEnv(gym.Env):\n",
    "    \"\"\"Custom Environment that follows gym interface\"\"\"\n",
    "    metadata = {'render.modes': ['human']}\n",
    "\n",
    "    def __init__(self, session):\n",
    "        super(EpiEnv, self).__init__()\n",
    "        self.epi = construct_epidemic(session)\n",
    "        total_population = np.sum(self.epi.static.default_state.obs.current_comp)\n",
    "        obs_count = self.epi.static.compartment_count * self.epi.static.locale_count * self.epi.static.group_count\n",
    "        action_count = 0\n",
    "        action_param_count =  0\n",
    "        for itv in self.epi.static.interventions:\n",
    "            if not itv.is_cost:\n",
    "                action_count += 1\n",
    "                action_param_count += len(itv.cp_list)\n",
    "        self.act_domain = np.zeros((action_param_count, 2), dtype=np.float64)\n",
    "        index = 0\n",
    "        for itv in self.epi.static.interventions:\n",
    "            if not itv.is_cost:\n",
    "                for cp in itv.cp_list:\n",
    "                    self.act_domain[index, 0] = cp.min_value\n",
    "                    self.act_domain[index, 1] = cp.max_value\n",
    "                    index += 1\n",
    "        # Define action and observation space\n",
    "        # They must be gym.spaces objects\n",
    "        # Example when using discrete actions:\n",
    "        self.action_space = spaces.Box(low=0, high=1, shape=(action_count,), dtype=np.float64)\n",
    "        # Example for using image as input:\n",
    "        self.observation_space = spaces.Box(low=0, high=total_population, shape=(obs_count,), dtype=np.float64)\n",
    "\n",
    "    def step(self, action):\n",
    "        expanded_action = np.zeros(len(self.act_domain), dtype=np.float64)\n",
    "        index = 0\n",
    "        for i in range(len(self.act_domain)):\n",
    "            if self.act_domain[i, 0] == self.act_domain[i, 1]:\n",
    "                expanded_action[i] = self.act_domain[i, 0]\n",
    "            else:\n",
    "                expanded_action[i] = action[index]\n",
    "                index += 1\n",
    "\n",
    "        epi_action = []\n",
    "        index = 0\n",
    "        for itv_id, itv in enumerate(self.epi.static.interventions):\n",
    "            if not itv.is_cost:\n",
    "                epi_action.append(construct_act(itv_id, expanded_action[index:index+len(itv.cp_list)]))\n",
    "                index += len(itv.cp_list)\n",
    "\n",
    "        total_r = 0\n",
    "        for i in range(PERIOD):\n",
    "            state, r, done = self.epi.step(epi_action)\n",
    "            total_r += r\n",
    "            if done:\n",
    "                break\n",
    "        return state.obs.current_comp.flatten(), total_r, done, dict()\n",
    "\n",
    "    def reset(self):\n",
    "        state = self.epi.reset()\n",
    "        return state.obs.current_comp.flatten()  # reward, done, info can't be included\n",
    "    def render(self, mode='human'):\n",
    "        pass\n",
    "    def close(self):\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_args(main_args = None):\n",
    "    # fmt: off\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--exp-name\", type=str, default=\"PPO\",\n",
    "        help=\"the name of this experiment\")\n",
    "    parser.add_argument(\"--gym-id\", type=str, default=\"HalfCheetahBulletEnv-v0\",\n",
    "        help=\"the id of the gym environment\")\n",
    "    parser.add_argument(\"--learning-rate\", type=float, default=3e-4,\n",
    "        help=\"the learning rate of the optimizer\")\n",
    "    parser.add_argument(\"--seed\", type=int, default=1,\n",
    "        help=\"seed of the experiment\")\n",
    "    parser.add_argument(\"--total-timesteps\", type=int, default=700000,\n",
    "        help=\"total timesteps of the experiments\")\n",
    "    parser.add_argument(\"--torch-deterministic\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "        help=\"if toggled, `torch.backends.cudnn.deterministic=False`\")\n",
    "    parser.add_argument(\"--cuda\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "        help=\"if toggled, cuda will be enabled by default\")\n",
    "    parser.add_argument(\"--track\", type=lambda x: bool(strtobool(x)), default=False, nargs=\"?\", const=True,\n",
    "        help=\"if toggled, this experiment will be tracked with Weights and Biases\")\n",
    "    parser.add_argument(\"--wandb-project-name\", type=str, default=\"ppo-implementation-details\",\n",
    "        help=\"the wandb's project name\")\n",
    "    parser.add_argument(\"--wandb-entity\", type=str, default=None,\n",
    "        help=\"the entity (team) of wandb's project\")\n",
    "    parser.add_argument(\"--capture-video\", type=lambda x: bool(strtobool(x)), default=False, nargs=\"?\", const=True,\n",
    "        help=\"weather to capture videos of the agent performances (check out `videos` folder)\")\n",
    "\n",
    "    parser.add_argument(\"--policy_plot_interval\", type=int, default=1,\n",
    "        help=\"seed of the experiment\")\n",
    "\n",
    "    # Algorithm specific arguments\n",
    "    parser.add_argument(\"--num-envs\", type=int, default=1,\n",
    "        help=\"the number of parallel game environments\")\n",
    "    parser.add_argument(\"--num-steps\", type=int, default=2048,\n",
    "        help=\"the number of steps to run in each environment per policy rollout\")\n",
    "    parser.add_argument(\"--anneal-lr\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "        help=\"Toggle learning rate annealing for policy and value networks\")\n",
    "    parser.add_argument(\"--gae\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "        help=\"Use GAE for advantage computation\")\n",
    "    parser.add_argument(\"--gamma\", type=float, default=0.99,\n",
    "        help=\"the discount factor gamma\")\n",
    "    parser.add_argument(\"--gae-lambda\", type=float, default=0.95,\n",
    "        help=\"the lambda for the general advantage estimation\")\n",
    "    parser.add_argument(\"--num-minibatches\", type=int, default=32,\n",
    "        help=\"the number of mini-batches\")\n",
    "    parser.add_argument(\"--update-epochs\", type=int, default=10,\n",
    "        help=\"the K epochs to update the policy\")\n",
    "    parser.add_argument(\"--norm-adv\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "        help=\"Toggles advantages normalization\")\n",
    "    parser.add_argument(\"--clip-coef\", type=float, default=0.2,\n",
    "        help=\"the surrogate clipping coefficient\")\n",
    "    parser.add_argument(\"--clip-vloss\", type=lambda x: bool(strtobool(x)), default=True, nargs=\"?\", const=True,\n",
    "        help=\"Toggles whether or not to use a clipped loss for the value function, as per the paper.\")\n",
    "    parser.add_argument(\"--ent-coef\", type=float, default=0.0,\n",
    "        help=\"coefficient of the entropy\")\n",
    "    parser.add_argument(\"--vf-coef\", type=float, default=0.5,\n",
    "        help=\"coefficient of the value function\")\n",
    "    parser.add_argument(\"--max-grad-norm\", type=float, default=0.5,\n",
    "        help=\"the maximum norm for the gradient clipping\")\n",
    "    parser.add_argument(\"--target-kl\", type=float, default=None,\n",
    "        help=\"the target KL divergence threshold\")\n",
    "    if main_args is not None:\n",
    "        args = parser.parse_args(main_args.split())\n",
    "    else:\n",
    "        args = parser.parse_args()\n",
    "    args.num_steps //= PERIOD\n",
    "    args.total_timesteps //= PERIOD\n",
    "    args.batch_size = int(args.num_envs * args.num_steps)\n",
    "    args.minibatch_size = int(args.batch_size // args.num_minibatches)\n",
    "    # fmt: on\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "epi_ids = [\"SIR_A\"]#, \"SIR_B\", \"SIRV_A\", \"SIRV_B\", \"COVID_A\", \"COVID_B\", \"COVID_C\"]\n",
    "\n",
    "def make_env(gym_id, seed, idx, capture_video, run_name):\n",
    "    def thunk():\n",
    "        if gym_id in epi_ids:\n",
    "            fp = open('jsons/{}.json'.format(gym_id), 'r')\n",
    "            session = json.load(fp)\n",
    "            env = EpiEnv(session)\n",
    "        else:\n",
    "            env = gym.make(gym_id)\n",
    "        env = gym.wrappers.RecordEpisodeStatistics(env)\n",
    "        if capture_video:\n",
    "            if idx == 0:\n",
    "                env = gym.wrappers.RecordVideo(env, f\"videos/{run_name}\")\n",
    "        env = gym.wrappers.ClipAction(env)\n",
    "        env = gym.wrappers.NormalizeObservation(env)\n",
    "        env = gym.wrappers.TransformObservation(env, lambda obs: np.clip(obs, -10, 10))\n",
    "        env = gym.wrappers.NormalizeReward(env)\n",
    "        env = gym.wrappers.TransformReward(env, lambda reward: np.clip(reward, -10, 10))\n",
    "        # Our env is deterministic\n",
    "        # env.seed(seed)\n",
    "        env.action_space.seed(seed)\n",
    "        env.observation_space.seed(seed)\n",
    "        return env\n",
    "\n",
    "    return thunk\n",
    "\n",
    "def make_primal_env(gym_id):\n",
    "    def thunk():\n",
    "        if gym_id in epi_ids:\n",
    "            fp = open('jsons/{}.json'.format(gym_id), 'r')\n",
    "            session = json.load(fp)\n",
    "            env = EpiEnv(session)\n",
    "        else:\n",
    "            env = gym.make(gym_id)\n",
    "        return env\n",
    "    return thunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_init(layer, std=np.sqrt(2), bias_const=0.0):\n",
    "    torch.nn.init.orthogonal_(layer.weight, std)\n",
    "    torch.nn.init.constant_(layer.bias, bias_const)\n",
    "    return layer\n",
    "\n",
    "\n",
    "class Agent(nn.Module):\n",
    "    def __init__(self, env):\n",
    "        super(Agent, self).__init__()\n",
    "        self.actor_mean_sigma = nn.Sequential(\n",
    "            layer_init(nn.Linear(np.array(env.observation_space.shape).prod(), 64)),\n",
    "            nn.Tanh(),\n",
    "            layer_init(nn.Linear(64, 64)),\n",
    "            nn.Tanh(),\n",
    "            layer_init(nn.Linear(64, 2*np.prod(env.action_space.shape)), std=0.01),\n",
    "        )\n",
    "        self.m = MultivariateNormal(\n",
    "            torch.zeros(np.prod(env.action_space.shape)), \n",
    "            torch.eye(np.prod(env.action_space.shape))\n",
    "        )\n",
    "\n",
    "    def get_action(self, x):\n",
    "        actor_mean_sigma = self.actor_mean_sigma(x)\n",
    "        action_mean = actor_mean_sigma[0,:np.prod(env.action_space.shape)]\n",
    "        action_sigma = actor_mean_sigma[0,np.prod(env.action_space.shape):]\n",
    "        epsilon = self.m.sample()\n",
    "        \n",
    "        action = action_mean + action_sigma * epsilon \n",
    "\t\t\n",
    "        # Apply the sigmoid to ensure the action is between 0 and 1\n",
    "        action = torch.sigmoid(action)\n",
    "        return action\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make_primal_env(args.gym_id)()\n",
    "agent = Agent(env)\n",
    "optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### FROM SCRATCH IMPLEMENTATION OF SVG\n",
    "NUM_SAMPLES = 10\n",
    "UNROLL_HORIZON = 10\n",
    "GAMMA = 0.99\n",
    "NUM_UPDATES\n",
    "# we need to store the rewards as we unroll\n",
    "\n",
    "rewards = torch.zeros((UNROLL_HORIZON, NUM_SAMPLES), requires_grad=True)\n",
    "actions = torch.zeros((UNROLL_HORIZON, NUM_SAMPLES) + env.action_space.shape)\n",
    "obs = torch.zeros((UNROLL_HORIZON, NUM_SAMPLES) + env.observation_space.shape)\n",
    "\n",
    "for update in range(0, NUM_UPDATES):\n",
    "\tfor i in range(NUM_SAMPLES):\n",
    "\t\tnext_obs = torch.Tensor(env.reset()).unsqueeze(0) # initial state\n",
    "\t\tfor t in range(UNROLL_HORIZON):\n",
    "\n",
    "\t\t\tobs[t,i] = next_obs\n",
    "\t\t\t# ALGO LOGIC: action logic\n",
    "\t\t\twith torch.no_grad():\n",
    "\t\t\t\taction = agent.get_action(next_obs)\n",
    "\t\t\tactions[t,i] = action\n",
    "\t\t\t\n",
    "\t\t\tnext_obs, reward, done, info = env.step(action.cpu().numpy())\n",
    "\t\t\twith torch.no_grad():\n",
    "\t\t\t\trewards[t,i] = torch.tensor(reward, dtype=torch.float32)\n",
    "\t\t\tnext_obs = torch.Tensor(next_obs).unsqueeze(0)\n",
    "\n",
    "\t# compute policy returns\n",
    "\tloss = torch.zeros((NUM_SAMPLES))\n",
    "\tfor t in range(UNROLL_HORIZON):\n",
    "\t\tloss +=  GAMMA**t * rewards[t,:]\n",
    "\n",
    "\tmean_loss = loss.mean()\n",
    "\n",
    "\toptimizer.zero_grad()\n",
    "\tmean_loss.backward()\n",
    "\toptimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00],\n",
       "         [2.2244e+06, 1.0000e+02, 0.0000e+00]],\n",
       "\n",
       "        [[2.1923e+06, 1.2807e+02, 3.2132e+04],\n",
       "         [2.1922e+06, 1.2801e+02, 3.2204e+04],\n",
       "         [2.1922e+06, 1.2823e+02, 3.2242e+04],\n",
       "         [2.1922e+06, 1.2809e+02, 3.2182e+04],\n",
       "         [2.1922e+06, 1.2805e+02, 3.2165e+04],\n",
       "         [2.1922e+06, 1.2808e+02, 3.2177e+04],\n",
       "         [2.1922e+06, 1.2796e+02, 3.2218e+04],\n",
       "         [2.1923e+06, 1.2792e+02, 3.2134e+04],\n",
       "         [2.1922e+06, 1.2807e+02, 3.2165e+04],\n",
       "         [2.1921e+06, 1.2785e+02, 3.2264e+04]],\n",
       "\n",
       "        [[2.1601e+06, 1.5714e+02, 6.4297e+04],\n",
       "         [2.1599e+06, 1.5745e+02, 6.4423e+04],\n",
       "         [2.1599e+06, 1.5770e+02, 6.4478e+04],\n",
       "         [2.1599e+06, 1.5747e+02, 6.4424e+04],\n",
       "         [2.1600e+06, 1.5718e+02, 6.4389e+04],\n",
       "         [2.1599e+06, 1.5742e+02, 6.4454e+04],\n",
       "         [2.1600e+06, 1.5722e+02, 6.4371e+04],\n",
       "         [2.1600e+06, 1.5701e+02, 6.4328e+04],\n",
       "         [2.1600e+06, 1.5708e+02, 6.4381e+04],\n",
       "         [2.1599e+06, 1.5724e+02, 6.4455e+04]],\n",
       "\n",
       "        [[2.1278e+06, 1.8995e+02, 9.6508e+04],\n",
       "         [2.1276e+06, 1.9004e+02, 9.6702e+04],\n",
       "         [2.1276e+06, 1.9075e+02, 9.6725e+04],\n",
       "         [2.1277e+06, 1.9047e+02, 9.6605e+04],\n",
       "         [2.1277e+06, 1.8998e+02, 9.6623e+04],\n",
       "         [2.1276e+06, 1.9052e+02, 9.6693e+04],\n",
       "         [2.1277e+06, 1.8951e+02, 9.6609e+04],\n",
       "         [2.1277e+06, 1.8973e+02, 9.6588e+04],\n",
       "         [2.1278e+06, 1.8984e+02, 9.6540e+04],\n",
       "         [2.1276e+06, 1.9015e+02, 9.6745e+04]],\n",
       "\n",
       "        [[2.0955e+06, 2.2594e+02, 1.2876e+05],\n",
       "         [2.0954e+06, 2.2598e+02, 1.2893e+05],\n",
       "         [2.0952e+06, 2.2674e+02, 1.2905e+05],\n",
       "         [2.0954e+06, 2.2659e+02, 1.2886e+05],\n",
       "         [2.0954e+06, 2.2597e+02, 1.2888e+05],\n",
       "         [2.0953e+06, 2.2670e+02, 1.2897e+05],\n",
       "         [2.0954e+06, 2.2537e+02, 1.2890e+05],\n",
       "         [2.0954e+06, 2.2567e+02, 1.2888e+05],\n",
       "         [2.0955e+06, 2.2574e+02, 1.2877e+05],\n",
       "         [2.0953e+06, 2.2599e+02, 1.2899e+05]],\n",
       "\n",
       "        [[2.0632e+06, 2.6417e+02, 1.6107e+05],\n",
       "         [2.0630e+06, 2.6420e+02, 1.6127e+05],\n",
       "         [2.0629e+06, 2.6530e+02, 1.6139e+05],\n",
       "         [2.0631e+06, 2.6534e+02, 1.6115e+05],\n",
       "         [2.0631e+06, 2.6437e+02, 1.6119e+05],\n",
       "         [2.0630e+06, 2.6536e+02, 1.6130e+05],\n",
       "         [2.0630e+06, 2.6393e+02, 1.6122e+05],\n",
       "         [2.0631e+06, 2.6448e+02, 1.6119e+05],\n",
       "         [2.0632e+06, 2.6386e+02, 1.6105e+05],\n",
       "         [2.0630e+06, 2.6420e+02, 1.6129e+05]],\n",
       "\n",
       "        [[2.0308e+06, 3.0419e+02, 1.9342e+05],\n",
       "         [2.0306e+06, 3.0374e+02, 1.9362e+05],\n",
       "         [2.0305e+06, 3.0521e+02, 1.9372e+05],\n",
       "         [2.0307e+06, 3.0543e+02, 1.9352e+05],\n",
       "         [2.0307e+06, 3.0403e+02, 1.9356e+05],\n",
       "         [2.0305e+06, 3.0594e+02, 1.9370e+05],\n",
       "         [2.0306e+06, 3.0400e+02, 1.9360e+05],\n",
       "         [2.0307e+06, 3.0420e+02, 1.9354e+05],\n",
       "         [2.0308e+06, 3.0354e+02, 1.9341e+05],\n",
       "         [2.0306e+06, 3.0437e+02, 1.9366e+05]],\n",
       "\n",
       "        [[1.9983e+06, 3.4415e+02, 2.2585e+05],\n",
       "         [1.9981e+06, 3.4414e+02, 2.2607e+05],\n",
       "         [1.9980e+06, 3.4512e+02, 2.2618e+05],\n",
       "         [1.9983e+06, 3.4611e+02, 2.2590e+05],\n",
       "         [1.9982e+06, 3.4456e+02, 2.2596e+05],\n",
       "         [1.9980e+06, 3.4623e+02, 2.2614e+05],\n",
       "         [1.9981e+06, 3.4375e+02, 2.2604e+05],\n",
       "         [1.9983e+06, 3.4466e+02, 2.2590e+05],\n",
       "         [1.9984e+06, 3.4395e+02, 2.2582e+05],\n",
       "         [1.9981e+06, 3.4520e+02, 2.2611e+05]],\n",
       "\n",
       "        [[1.9658e+06, 3.8354e+02, 2.5832e+05],\n",
       "         [1.9656e+06, 3.8266e+02, 2.5855e+05],\n",
       "         [1.9655e+06, 3.8399e+02, 2.5866e+05],\n",
       "         [1.9658e+06, 3.8628e+02, 2.5838e+05],\n",
       "         [1.9657e+06, 3.8407e+02, 2.5843e+05],\n",
       "         [1.9655e+06, 3.8643e+02, 2.5861e+05],\n",
       "         [1.9656e+06, 3.8267e+02, 2.5852e+05],\n",
       "         [1.9658e+06, 3.8405e+02, 2.5838e+05],\n",
       "         [1.9658e+06, 3.8265e+02, 2.5830e+05],\n",
       "         [1.9656e+06, 3.8472e+02, 2.5858e+05]],\n",
       "\n",
       "        [[1.9333e+06, 4.2023e+02, 2.9084e+05],\n",
       "         [1.9330e+06, 4.1941e+02, 2.9106e+05],\n",
       "         [1.9329e+06, 4.2033e+02, 2.9117e+05],\n",
       "         [1.9332e+06, 4.2320e+02, 2.9090e+05],\n",
       "         [1.9332e+06, 4.2121e+02, 2.9094e+05],\n",
       "         [1.9330e+06, 4.2297e+02, 2.9113e+05],\n",
       "         [1.9331e+06, 4.1910e+02, 2.9103e+05],\n",
       "         [1.9332e+06, 4.2087e+02, 2.9089e+05],\n",
       "         [1.9333e+06, 4.1946e+02, 2.9081e+05],\n",
       "         [1.9330e+06, 4.2139e+02, 2.9109e+05]]])"
      ]
     },
     "execution_count": 160,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.4974, 0.4975],\n",
       "         [0.4985, 0.4982],\n",
       "         [0.4991, 0.4957],\n",
       "         [0.4982, 0.4973],\n",
       "         [0.4979, 0.4978],\n",
       "         [0.4981, 0.4975],\n",
       "         [0.4987, 0.4988],\n",
       "         [0.4974, 0.4992],\n",
       "         [0.4979, 0.4976],\n",
       "         [0.4994, 0.5000]],\n",
       "\n",
       "        [[0.4975, 0.4987],\n",
       "         [0.4984, 0.4949],\n",
       "         [0.4986, 0.4952],\n",
       "         [0.4987, 0.4957],\n",
       "         [0.4984, 0.4980],\n",
       "         [0.4993, 0.4960],\n",
       "         [0.4973, 0.4964],\n",
       "         [0.4980, 0.4981],\n",
       "         [0.4983, 0.4991],\n",
       "         [0.4979, 0.4948]],\n",
       "\n",
       "        [[0.4978, 0.4970],\n",
       "         [0.4989, 0.4994],\n",
       "         [0.4984, 0.4961],\n",
       "         [0.4973, 0.4961],\n",
       "         [0.4982, 0.4970],\n",
       "         [0.4982, 0.4952],\n",
       "         [0.4982, 0.5014],\n",
       "         [0.4986, 0.4975],\n",
       "         [0.4970, 0.4972],\n",
       "         [0.4990, 0.4966]],\n",
       "\n",
       "        [[0.4980, 0.4973],\n",
       "         [0.4977, 0.4973],\n",
       "         [0.4992, 0.4982],\n",
       "         [0.4980, 0.4970],\n",
       "         [0.4981, 0.4973],\n",
       "         [0.4983, 0.4967],\n",
       "         [0.4986, 0.4971],\n",
       "         [0.4986, 0.4972],\n",
       "         [0.4976, 0.4977],\n",
       "         [0.4978, 0.4983]],\n",
       "\n",
       "        [[0.4984, 0.4990],\n",
       "         [0.4987, 0.4990],\n",
       "         [0.4988, 0.4975],\n",
       "         [0.4981, 0.4965],\n",
       "         [0.4984, 0.4979],\n",
       "         [0.4988, 0.4971],\n",
       "         [0.4986, 0.4963],\n",
       "         [0.4984, 0.4952],\n",
       "         [0.4979, 0.4994],\n",
       "         [0.4983, 0.4989]],\n",
       "\n",
       "        [[0.4984, 0.4976],\n",
       "         [0.4986, 0.5000],\n",
       "         [0.4982, 0.4989],\n",
       "         [0.4989, 0.4983],\n",
       "         [0.4986, 0.4997],\n",
       "         [0.4992, 0.4956],\n",
       "         [0.4989, 0.4973],\n",
       "         [0.4984, 0.4996],\n",
       "         [0.4987, 0.4992],\n",
       "         [0.4988, 0.4967]],\n",
       "\n",
       "        [[0.4991, 0.5002],\n",
       "         [0.4994, 0.4976],\n",
       "         [0.4995, 0.5007],\n",
       "         [0.4983, 0.4975],\n",
       "         [0.4987, 0.4972],\n",
       "         [0.4993, 0.4997],\n",
       "         [0.4994, 0.5010],\n",
       "         [0.4981, 0.4977],\n",
       "         [0.4989, 0.4976],\n",
       "         [0.4994, 0.4961]],\n",
       "\n",
       "        [[0.4993, 0.4983],\n",
       "         [0.4994, 0.5020],\n",
       "         [0.4994, 0.5007],\n",
       "         [0.4993, 0.4961],\n",
       "         [0.4993, 0.4981],\n",
       "         [0.4993, 0.4957],\n",
       "         [0.4994, 0.4999],\n",
       "         [0.4994, 0.4987],\n",
       "         [0.4995, 0.5014],\n",
       "         [0.4993, 0.4983]],\n",
       "\n",
       "        [[0.4994, 0.4999],\n",
       "         [0.4993, 0.4988],\n",
       "         [0.4994, 0.5010],\n",
       "         [0.4994, 0.5002],\n",
       "         [0.4994, 0.4983],\n",
       "         [0.4994, 0.5015],\n",
       "         [0.4994, 0.5003],\n",
       "         [0.4993, 0.4995],\n",
       "         [0.4994, 0.4988],\n",
       "         [0.4993, 0.5002]],\n",
       "\n",
       "        [[0.4995, 0.5004],\n",
       "         [0.4994, 0.4998],\n",
       "         [0.4995, 0.4999],\n",
       "         [0.4994, 0.4995],\n",
       "         [0.4995, 0.4993],\n",
       "         [0.4994, 0.4999],\n",
       "         [0.4994, 0.4997],\n",
       "         [0.4994, 0.5001],\n",
       "         [0.4995, 0.4999],\n",
       "         [0.4995, 0.4998]]])"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-19045084., -19052208., -19063328., -19053060., -19048184., -19062802.,\n",
       "        -19049478., -19046488., -19042106., -19055388.],\n",
       "       grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 162,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-19051812., grad_fn=<MeanBackward0>)"
      ]
     },
     "execution_count": 163,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean_loss = loss.mean()\n",
    "mean_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-1805298.8750, -1808701.6250, -1808317.3750, -1807128.0000,\n",
       "         -1806821.0000, -1807047.0000, -1809675.1250, -1806623.8750,\n",
       "         -1806650.7500, -1812439.2500],\n",
       "        [-1839678.2500, -1838968.5000, -1840207.1250, -1840628.0000,\n",
       "         -1841508.3750, -1842238.8750, -1837413.3750, -1840244.6250,\n",
       "         -1842057.8750, -1837561.5000],\n",
       "        [-1876282.2500, -1881208.1250, -1877834.8750, -1874911.7500,\n",
       "         -1877283.5000, -1876487.0000, -1880721.0000, -1878457.6250,\n",
       "         -1874277.2500, -1879327.8750],\n",
       "        [-1918479.6250, -1917720.0000, -1923142.5000, -1918874.8750,\n",
       "         -1918662.5000, -1919500.1250, -1919224.8750, -1919669.7500,\n",
       "         -1917587.2500, -1919044.2500],\n",
       "        [-1965612.8750, -1966545.0000, -1966723.5000, -1963988.3750,\n",
       "         -1965007.0000, -1966199.1250, -1963422.6250, -1962641.7500,\n",
       "         -1964466.6250, -1965289.2500],\n",
       "        [-2011988.3750, -2013915.6250, -2013621.0000, -2015078.3750,\n",
       "         -2014179.3750, -2013984.2500, -2012554.0000, -2013626.8750,\n",
       "         -2013193.7500, -2012267.2500],\n",
       "        [-2064336.6250, -2062721.1250, -2066979.8750, -2061936.6250,\n",
       "         -2060939.0000, -2066779.8750, -2065216.1250, -2060077.3750,\n",
       "         -2061128.5000, -2062454.6250],\n",
       "        [-2111565.7500, -2114287.2500, -2114513.5000, -2112663.5000,\n",
       "         -2111916.0000, -2112301.0000, -2112290.0000, -2112716.2500,\n",
       "         -2113807.0000, -2112964.7500],\n",
       "        [-2159616.0000, -2157456.7500, -2160695.2500, -2163210.5000,\n",
       "         -2159071.0000, -2164161.0000, -2158632.7500, -2159732.5000,\n",
       "         -2157519.5000, -2161017.5000],\n",
       "        [-2202580.5000, -2201029.5000, -2202366.2500, -2205545.0000,\n",
       "         -2203108.7500, -2205430.0000, -2200564.2500, -2202973.5000,\n",
       "         -2201449.7500, -2203619.0000]], requires_grad=True)"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rewards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running SIR_A__PPO__1__1734145470\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/342 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "global_step=53, episodic_return=-2089364992.0\n",
      "global_step=106, episodic_return=-2093162112.0\n",
      "global_step=159, episodic_return=-2092365824.0\n",
      "global_step=212, episodic_return=-2094185600.0\n",
      "global_step=265, episodic_return=-2092165888.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/10 [00:00<?, ?it/s]\n",
      "  0%|          | 0/342 [00:29<?, ?it/s]\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "element 0 of tensors does not require grad and does not have a grad_fn",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[25], line 114\u001b[0m\n\u001b[0;32m    111\u001b[0m loss \u001b[38;5;241m=\u001b[39m pg_loss \n\u001b[0;32m    113\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[1;32m--> 114\u001b[0m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    115\u001b[0m nn\u001b[38;5;241m.\u001b[39mutils\u001b[38;5;241m.\u001b[39mclip_grad_norm_(agent\u001b[38;5;241m.\u001b[39mparameters(), args\u001b[38;5;241m.\u001b[39mmax_grad_norm)\n\u001b[0;32m    116\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n",
      "File \u001b[1;32mc:\\Users\\IBM\\Projects\\rl\\rl-mai-paper-code\\RL-Epidemic-Benchmark\\.venv\\lib\\site-packages\\torch\\_tensor.py:581\u001b[0m, in \u001b[0;36mTensor.backward\u001b[1;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[0;32m    571\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m    572\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[0;32m    573\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[0;32m    574\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    579\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[0;32m    580\u001b[0m     )\n\u001b[1;32m--> 581\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    582\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[0;32m    583\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\IBM\\Projects\\rl\\rl-mai-paper-code\\RL-Epidemic-Benchmark\\.venv\\lib\\site-packages\\torch\\autograd\\__init__.py:347\u001b[0m, in \u001b[0;36mbackward\u001b[1;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[0;32m    342\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[0;32m    344\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[0;32m    345\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[0;32m    346\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[1;32m--> 347\u001b[0m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    348\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    349\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    350\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    351\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    352\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    353\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m    354\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[0;32m    355\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\IBM\\Projects\\rl\\rl-mai-paper-code\\RL-Epidemic-Benchmark\\.venv\\lib\\site-packages\\torch\\autograd\\graph.py:825\u001b[0m, in \u001b[0;36m_engine_run_backward\u001b[1;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[0;32m    823\u001b[0m     unregister_hooks \u001b[38;5;241m=\u001b[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[0;32m    824\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 825\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m Variable\u001b[38;5;241m.\u001b[39m_execution_engine\u001b[38;5;241m.\u001b[39mrun_backward(  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[0;32m    826\u001b[0m         t_outputs, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[0;32m    827\u001b[0m     )  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[0;32m    828\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m    829\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
      "\u001b[1;31mRuntimeError\u001b[0m: element 0 of tensors does not require grad and does not have a grad_fn"
     ]
    }
   ],
   "source": [
    "# from tqdm import tqdm\n",
    "# args = parse_args(\"--gym-id SIR_A --seed 1\")\n",
    "# run_name = f\"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}\"\n",
    "# print(\"Running\", run_name)\n",
    "# writer = SummaryWriter(f\"runs/{run_name}\")\n",
    "# writer.add_text(\n",
    "#     \"hyperparameters\",\n",
    "#     \"|param|value|\\n|-|-|\\n%s\" % (\"\\n\".join([f\"|{key}|{value}|\" for key, value in vars(args).items()])),\n",
    "# )\n",
    "\n",
    "# # TRY NOT TO MODIFY: seeding\n",
    "# random.seed(args.seed)\n",
    "# np.random.seed(args.seed)\n",
    "# torch.manual_seed(args.seed)\n",
    "# torch.backends.cudnn.deterministic = args.torch_deterministic\n",
    "\n",
    "# device = torch.device(\"cuda\" if torch.cuda.is_available() and args.cuda else \"cpu\")\n",
    "\n",
    "# # env setup\n",
    "# envs = gym.vector.SyncVectorEnv(\n",
    "#     [make_env(args.gym_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]\n",
    "# )\n",
    "# # envs = gym.vector.AsyncVectorEnv(\n",
    "# #     [make_env(args.gym_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]\n",
    "# # )\n",
    "# test_env = make_primal_env(args.gym_id)()\n",
    "# # test_env = make_env(args.gym_id, args.seed, 0, args.capture_video, run_name)()\n",
    "# assert isinstance(envs.single_action_space, gym.spaces.Box), \"only continuous action space is supported\"\n",
    "\n",
    "# agent = Agent(envs).to(device)\n",
    "# optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)\n",
    "\n",
    "# # ALGO Logic: Storage setup\n",
    "# obs = torch.zeros((args.num_steps, args.num_envs) + envs.single_observation_space.shape).to(device)\n",
    "# actions = torch.zeros((args.num_steps, args.num_envs) + envs.single_action_space.shape).to(device)\n",
    "# logprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)\n",
    "# rewards = torch.zeros((args.num_steps, args.num_envs)).to(device)\n",
    "# dones = torch.zeros((args.num_steps, args.num_envs)).to(device)\n",
    "# values = torch.zeros((args.num_steps, args.num_envs)).to(device)\n",
    "\n",
    "# # TRY NOT TO MODIFY: start the game\n",
    "# global_step = 0\n",
    "# start_time = time.time()\n",
    "# next_obs = torch.Tensor(envs.reset()).to(device)\n",
    "# next_done = torch.zeros(args.num_envs).to(device)\n",
    "# num_updates = args.total_timesteps // args.batch_size\n",
    "# ROLLOUT_HORIZON = 35\n",
    "\n",
    "# csv_file = open('runs/{}/records.csv'.format(run_name), 'w')\n",
    "\n",
    "# for update in tqdm(range(1, num_updates + 1)):\n",
    "#     # Annealing the rate if instructed to do so.\n",
    "#     if args.anneal_lr:\n",
    "#         frac = 1.0 - (update - 1.0) / num_updates\n",
    "#         lrnow = frac * args.learning_rate\n",
    "#         optimizer.param_groups[0][\"lr\"] = lrnow\n",
    "\n",
    "#     for step in range(0, args.num_steps):\n",
    "#         global_step += 1 * args.num_envs\n",
    "#         obs[step] = next_obs\n",
    "#         dones[step] = next_done\n",
    "\n",
    "#         # ALGO LOGIC: action logic\n",
    "#         with torch.no_grad():\n",
    "#             action, _, _, _ = agent.get_action(next_obs)\n",
    "#         actions[step] = action\n",
    "\n",
    "#         # TRY NOT TO MODIFY: execute the game and log data.\n",
    "#         next_obs, reward, done, info = envs.step(action.cpu().numpy())\n",
    "#         rewards[step] = torch.tensor(reward).to(device).view(-1)\n",
    "#         next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(done).to(device)\n",
    "\n",
    "#         for item in info:\n",
    "#             if \"episode\" in item.keys():\n",
    "#                 print(f\"global_step={global_step}, episodic_return={item['episode']['r']}\")\n",
    "#                 writer.add_scalar(\"charts/episodic_return\", item[\"episode\"][\"r\"], global_step)\n",
    "#                 writer.add_scalar(\"charts/episodic_length\", item[\"episode\"][\"l\"], global_step)\n",
    "#                 break\n",
    "\n",
    "#     # bootstrap value if not done\n",
    "#     with torch.no_grad():\n",
    "#         returns = torch.zeros_like(rewards).to(device)\n",
    "#         for t in reversed(range(args.num_steps)):\n",
    "#             if t == args.num_steps - 1:\n",
    "#                 nextnonterminal = 1.0 - next_done\n",
    "#                 next_return = 0.0\n",
    "#             else:\n",
    "#                 nextnonterminal = 1.0 - dones[t + 1]\n",
    "#                 next_return = returns[t + 1]\n",
    "#             returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return\n",
    "\n",
    "#     # flatten the batch\n",
    "#     b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)\n",
    "#     b_logprobs = logprobs.reshape(-1)\n",
    "#     b_actions = actions.reshape((-1,) + envs.single_action_space.shape)\n",
    "#     b_returns = returns.reshape(-1)\n",
    "\n",
    "#     # Optimizing the policy and value network\n",
    "#     b_inds = np.arange(args.batch_size)\n",
    "#     clipfracs = []\n",
    "#     for epoch in tqdm(range(max(1, args.update_epochs))): # eventually remove epoch as we look at trajectories only once\n",
    "#         np.random.shuffle(b_inds)\n",
    "#         for start in range(0, args.batch_size, args.minibatch_size):\n",
    "#             # TODO: do we actually need mini-batches as we look at all trajectories at once?\n",
    "#             end = start + args.minibatch_size\n",
    "#             mb_inds = b_inds[start:end]\n",
    "\n",
    "#             mb_returns = b_returns[mb_inds]\n",
    "#             pg_loss = -mb_returns\n",
    "            \n",
    "#             loss = pg_loss \n",
    "\n",
    "#             optimizer.zero_grad()\n",
    "#             loss.backward()\n",
    "#             nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)\n",
    "#             optimizer.step()\n",
    "\n",
    "\n",
    "#     y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy()\n",
    "#     var_y = np.var(y_true)\n",
    "#     explained_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y\n",
    "\n",
    "#     # TRY NOT TO MODIFY: record rewards for plotting purposes\n",
    "#     writer.add_scalar(\"charts/learning_rate\", optimizer.param_groups[0][\"lr\"], global_step)\n",
    "#     writer.add_scalar(\"losses/value_loss\", v_loss.item(), global_step)\n",
    "#     writer.add_scalar(\"losses/policy_loss\", pg_loss.item(), global_step)\n",
    "#     writer.add_scalar(\"losses/entropy\", entropy_loss.item(), global_step)\n",
    "#     writer.add_scalar(\"losses/old_approx_kl\", old_approx_kl.item(), global_step)\n",
    "#     writer.add_scalar(\"losses/approx_kl\", approx_kl.item(), global_step)\n",
    "#     writer.add_scalar(\"losses/clipfrac\", np.mean(clipfracs), global_step)\n",
    "#     writer.add_scalar(\"losses/explained_variance\", explained_var, global_step)\n",
    "#     print(\"SPS:\", int(global_step / (time.time() - start_time)))\n",
    "#     writer.add_scalar(\"charts/SPS\", int(global_step / (time.time() - start_time)), global_step)\n",
    "\n",
    "#     # PLOT POLICY\n",
    "#     if args.gym_id in epi_ids and (update - 1) % args.policy_plot_interval == 0:\n",
    "#         test_obs = torch.Tensor(test_env.reset())\n",
    "#         env_obs = torch.Tensor(envs.reset()).to(device)\n",
    "#         timestep = 0\n",
    "#         total_r = 0\n",
    "#         done = False\n",
    "#         itv_line = []\n",
    "#         while not done:\n",
    "#             with torch.no_grad():\n",
    "#                 action_mean = agent.get_action_mean(env_obs)\n",
    "#                 test_action_mean = torch.mean(action_mean, 0)\n",
    "#                 test_action_mean = torch.clamp(test_action_mean, 0, 1)\n",
    "             \n",
    "#             test_obs, r, done, _ = test_env.step(test_action_mean.cpu().numpy())\n",
    "#             test_obs = torch.Tensor(test_obs)\n",
    "#             itv_index = 0\n",
    "#             itv_array = []\n",
    "#             for itv in test_env.epi.static.interventions:\n",
    "#                 if not itv.is_cost:\n",
    "#                     v = float(test_action_mean[itv_index])\n",
    "#                     writer.add_scalar('charts/policy_{}/{}'.format(global_step, itv.name), v, timestep)\n",
    "#                     itv_array.append(v)\n",
    "#                     itv_index += 1\n",
    "#             itv_line.append(itv_array)\n",
    "            \n",
    "#             env_obs, _, _, _ = envs.step(action_mean.cpu().numpy())\n",
    "#             env_obs = torch.Tensor(env_obs).to(device)\n",
    "\n",
    "#             total_r += r\n",
    "#             timestep += PERIOD\n",
    "            \n",
    "#         line = '|'.join([str(global_step), str(total_r), str(itv_line)]) + '\\n'\n",
    "#         csv_file.write(line)\n",
    "#         writer.add_scalar('charts/learning_curve', total_r, global_step)\n",
    "#         print(\"At global step {}, total_rewards={}\".format(global_step, total_r))\n",
    "\n",
    "# csv_file.close()\n",
    "# envs.close()\n",
    "# writer.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}